최대 1 분 소요

문제 링크

풀이

# 최소신장트리 찾기
# 집합의 union, find 알고리즘 이용
# 간선의 수가 N-1이어야하고, 사이클이 안만들어져야함
# 노드의 루트노드가 다를 경우에만 union 진행

def find(parent, i):
    if parent[i] == i:
        return i
    else:
        parent[i] = find(parent, parent[i])
    return parent[i]
    
def union(parent, rank, x, y):
    xroot = find(parent, x)
    yroot = find(parent, y)
    
    if rank[xroot] < rank[yroot]:
        parent[xroot] = yroot
    elif rank[xroot] > rank[yroot]:
        parent[yroot] = xroot
    else:
        parent[xroot] = yroot
        rank[yroot] += 1
    
def solution(n, costs):
    costs.sort(key = lambda x: x[2]) # 비용낮은 순으로 정렬
    parent = [i for i in range(n)] # 노드별 부모노드
    edges = 0 # 간선 수
    rank = [0] * n 
    min_cost = 0
    
    for cost in costs:
        if edges == n -1:
            break
        x = find(parent, cost[0]) # 노드의 루트노드 비교
        y = find(parent, cost[1])
        
        if x != y:
            union(parent, rank, x, y)
            min_cost += cost[2]
            edges += 1
            
    return min_cost
        

댓글남기기