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

링크

백준 1197 최소 스패닝 트리

싸피에서 Prim, Kruskal, Dijkstra를 배웠고 연습하기 위해 풀었다.

해당 문제의 경우 Prim과 Kruskal중 하나를 사용해 최소신장트리를 만드는 문제이다.

Prim알고리즘을 활용해 풀었다.
임의의 정점 하나를 선택해서 시작하며 해당 정점에서부터 트리를 늘려가는 방식이다.

• Prim
  • 임의의 정점 하나를 선택해서 시작
  • 선택한 정점과 인접하는 정점들 중 가중치가 최소인 정점을 선택
  • 모든 정점이 선택될 때까지 반복

분명 풀고나서 효율성을 고려해서 짰다고 생각했는데 시간이 5초씩 걸렸다.. 다른 사람들의 코드를 봤는데 상위권 사람들은 대부분 크루스칼로 풀었길래 아니 크루스칼이랑 이렇게 차이가 난다고?? 하고 있었는데 범인은 파이썬 input이었다..
input을 stdin으로 바꿔주니 속도가 많이 차이났다. 인풋이 많은 문제는 까먹지 말자

정답 코드

from heapq import heappush, heappop
import sys; input=sys.stdin.readline

def prim(start, weight):
    visit = [0] * (V + 1) # 정점 방문 처리
    q = [[weight, start]] # 힙 구조를 사용하기 위해 가중치를 앞에 둠
    ans = 0 # 가중치 합
    cnt = 0 # 간선의 개수
    while cnt < V: # 간선의 개수 최대는 V-1
        k, v = heappop(q)
        if visit[v]: continue # 이미 방문한 정점이면 지나감
        visit[v] = 1 # 방문안했으면 방문처리
        ans += k # 해당 정점까지의 가중치를 더해줌
        cnt += 1 # 간선의 갯수 더해줌
        for u, w in G[v]: # 해당 정점의 간선정보를 불러옴
            heappush(q, [w, u]) # 힙에 넣어줌
    return ans

V, E = map(int, input().split())
G = [[] for _ in range(V + 1)]
for _ in range(E):
    u, v, w = map(int, input().split())
    G[u].append([v, w])
    G[v].append([u, w])

print(prim(1, 0))

알게된 것👨‍💻

• 프림