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

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

1197번: 최소 스패닝 트리

백트래킹 문제를 풀다가 계속 틀려서 쉬운문제를 하나 가져왔다.

최소 신장 트리를 구하는 문제이고 prim 알고리즘을 사용해서 풀었다.

다익스트라랑 크게 차이가 없는듯 하다.

from heapq import heappop, heappush
def mst(g):
    ret = 0
    visit = [0 for _ in range(len(g))]
    h = [(0, 1)] # cost, start
    while h:
        cur_cost, cur = heappop(h)
        if visit[cur]: continue
        visit[cur] = 1
        ret += cur_cost
        for adj_c, adj in g[cur]:
            heappush(h, (adj_c, adj))
    return ret

def main():
    v, e = map(int, input().split())
    g = [[] for _ in range(v + 1)]
    for _ in range(e):
        a, b, c = map(int, input().split())
        g[a].append((c, b))
        g[b].append((c, a))

    print(mst(g))


if __name__ == "__main__":
    main()

 

백트래킹 문제를 많이 풀어야하느데....