[MST] 1197 최소 스패닝 트리

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

1197번: 최소 스패닝 트리

# 2023.01.14
import sys
input = sys.stdin.readline
import heapq

v, e = map(int, input().split()) # 정점, 간선
graph=[[] for _ in range(v+1)] 
for i in range(e): # 간선의 수만큼 반복
    a, b, c = map(int, input().split())
    graph[a].append([c, b])
    graph[b].append([c, a])


heap=[[0, 1]] # weight, node 순으로 push
visited=[False]*(v+1)
answer=0
cnt=0
while heap:
    if cnt==v:
        break
    weight, node = heapq.heappop(heap) # 노드와 가중치
    if not visited[node]:
        answer+=weight
        visited[node]=True # 방문
        cnt+=1
        for i in graph[node]:
            heapq.heappush(heap, i)

print(answer)

---

# 2023.01.18
import sys
input = sys.stdin.readline
import heapq

v, e = map(int, input().split())
graph=[[] for _ in range(v+1)] # 노드의 수
for i in range(e):
    a, b, c = map(int, input().split()) # 노드1, 노드2, 가중치
    graph[a].append([c, b]) 
    graph[b].append([c, a])

heap=[[0, 1]]
visited=[0]*(v+1)
answer=0
while heap:
    if sum(visited)==v: # 모든 노드를 연결하였다면? 
        break
    weight, node = heapq.heappop(heap)
    if not visited[node]: # 방문하지 않은 노드라면? 해당 노드를 연결
        answer+=weight
        visited[node]=1
        for i in graph[node]: # 현재 노드와 연결되어 있는 모든 노드를 minHeap에 넣어줌
            heapq.heappush(heap, i)

print(answer)

pypy3로 제출해야 한다(pypy3가 더 빠름)

if sum(visited)==v: 에서 sum이라는 함수가 시간이 많이 걸려서 python3로 제출하면 시간초과가 발생한다

 

python3로 시간초과가 나지 않게 하려면?

# 2023.01.18
import sys
input = sys.stdin.readline
import heapq

v, e = map(int, input().split())
graph=[[] for _ in range(v+1)] # 노드의 수
for i in range(e):
    a, b, c = map(int, input().split()) # 노드1, 노드2, 가중치
    graph[a].append([c, b]) 
    graph[b].append([c, a])

heap=[[0, 1]]
visited=[False]*(v+1)
connectedNode=0
answer=0
while heap:
    if connectedNode==v: # 모든 노드를 연결하였다면?
        break
    weight, node = heapq.heappop(heap)
    if not visited[node]: # 방문하지 않은 노드라면?
        answer+=weight
        visited[node]=True
        connectedNode+=1
        for i in graph[node]:
            heapq.heappush(heap, i)

print(answer)

별도의 connectedNode라는 변수를 설정하고 노드를 연결할때마다 +1을 해주어서 모든 노드가 연결이 되었는지 check해주어야 한다!