[BOJ] 1197 최소 스패닝 트리

백준 골드4 1197 최소 스패닝 트리

문제 링크

문제 설명

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

코드

import sys
input = sys.stdin.readline

# V: 정점의 개수, E: 간선의 개수
V, E = map(int, input().split())

# 인접간선 행렬 만들기
edges = {}
for _ in range(E):
    a,b,c = map(int, input().split())
    if a in edges.keys():
        edges[a].append((b,c))
    else:
        edges[a] = [(b,c)]
    if b in edges.keys():
        edges[b].append((a,c))
    else:
        edges[b] = [(a,c)]

INF = sys.maxsize    

# 방문처리용
visit = [False]*(V+1)

# 가중치를 저장할 용
dist = [INF]*(V+1)

# 1번을 시작점으로 잡자
dist[1] = 0

ans = 0
for _ in range(V-1):
    mini = INF
    idx = 0
    # 방문 안한 정점 중 가장 작은 값
    for i in range(1, V+1):
        if not visit[i] and dist[i] < mini:
            mini = dist[i]
            idx = i
    
    # 뽑은 idx에 방문처리
    visit[idx] = True
    
    if not idx in edges.keys():
        continue
    
    # 다음 갱신
    for i in edges[idx]:
        if not visit[i[0]] and dist[i[0]] > i[1]:
            dist[i[0]] = i[1]

for i in range(1, V+1):
    ans += dist[i]

print(ans)

풀이

처음엔 간선의 정보를 딕셔너리가 아닌 인접행렬로 받았는데 메모리초과가 났다.
만약에 pypy로 안내고 python으로 제출했으면 괜찮았을까...?
흐음.. 암튼 최소신장트리로 풀어야겠다고 생각을 했고
프림 알고리즘 중 우선순위 큐가 아닌 반복문을 사용하여 풀었다.
이것 저것 해보고 MST중 가장 손에 익는 방법을 찾아봐야 겠다.