[python] 백준 6497번: 전력난

문제

https://www.acmicpc.net/problem/6497

알고리즘

최소 신장 트리, Union

아이디어

union 알고리즘인 find_parent, union_parent 를 기반으로,
heapq 를 통해 거리가 최소인 순서대로 cycle이 존재하는지를 판단해 cycle 이 없는 경우엔 최적의 해임이 보장되므로

이를 추가한다.

코드

from sys import stdin
import heapq


def find_parent(parent,a):
    if parent[a] != a:
        parent[a] = find_parent(parent,parent[a])
    return parent[a]


def union_parent(parent,a, b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)

    if a > b:
        parent[a] = b
    else:
        parent[b] = a

while True:
    n, m = map(int, stdin.readline().split())
    if (n,m) == (0,0):
        break
    parent = [0] * n

    for i in range(n):
        parent[i] = i

    q = []
    heapq.heapify(q)
    tot = 0
    ans = 0

    for _ in range(m):
        a, b, c = map(int, stdin.readline().split())
        heapq.heappush(q, (c, a, b))


    while q:
        dist, s, e = heapq.heappop(q)
        tot += dist
        if find_parent(parent,s) != find_parent(parent,e):
            union_parent(parent,s, e)
            ans += dist

    print(tot - ans)

참고사항: heapq 를 쓰지 않고 list 를 써도 무방하다
참고사항2: 여러 테스트 케이스가 주어지는 경우인데 그걸 못봐서 계속 틀렸다