[백준 1197] 최소 스패닝 트리

Junyoung Park·2022년 3월 1일
0

코딩테스트

목록 보기
138/631
post-thumbnail

1. 문제 설명

최소 스패닝 트리

2. 문제 분석

크루스칼 알고리즘을 통해 최소 신장 트리를 구할 수 있다. 이때 Union-Find를 통해 각 노드의 루트 노드를 구할 수 있다.

  • Find 함수를 재귀 또는 반복문으로 풀었을 때에는 시간 초과가 났는데, 메모라이제이션을 통해 효율적으로 풀 수 있었다.

3. 나의 풀이

import heapq
import sys

v, e = map(int, sys.stdin.readline().rstrip().split())
parents = [i for i in range(v+1)]
heap = []

for _ in range(e):
    tail, head, cost = map(int, sys.stdin.readline().rstrip().split())
    heapq.heappush(heap, [cost, tail, head])
    # 최소 신장 트리이므로 최소 비용 순서대로 간선을 정렬한다.
total = 0

def find(node):
    if parents[node] == node: return node
    # 루트 노드가 자기 자신이면 그대로 리턴한다. 
    parents[node] = find(parents[node])
    return parents[node]
    # 루트 노드를 찾아 리턴한다. (재귀 효율화)

def union(node1, node2):
    root1 = find(node1)
    root2 = find(node2)

    if root1 == root2: return False
    else:
        parents[root1] = root2
        return True
    # 루트 노드가 서로 같다면 이미 연결한 노드이므로 신장 트리에 더할 수 없다.

while heap:
    cost, tail, head = heapq.heappop(heap)

    if union(tail, head):
        total += cost
        # 사이클을 만들지 않는 간선이라면 더하자.
print(total)
profile
JUST DO IT

0개의 댓글