본문 바로가기

코딩테스트/백준

[백준] 1197번 : 최소 스패닝 트리 (Python)

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

 

주어진 그래프의 최소 스패닝 트리를 구하는 문제다.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는

가중치의 합이 최소인 트리를 말한다.

 

import sys
import heapq

sys.setrecursionlimit(10**6)  # 재귀 깊이 제한 1,000,000으로 늘림
input = sys.stdin.readline

V, E = map(int, input().strip().split())  # 노드 수, 에지 수

pq = []
parent = [i for i in range(V + 1)]  # 대표 노드 저장 리스트

# 우선순위 큐에 에지 정보 저장
for i in range(E):
    s, e, v = map(int, input().strip().split())
    heapq.heappush(pq, (v, s, e))  # 맨 앞의 값이 정렬의 기준, 그러니 가중치부터 넣어준다.

# find 연산
def find(a):
    if parent[a] != a # 경로 압축, 한번 find 연산을 할 때 중간경로도 다 대표노드로 바꾸게 된다. 
        parent[a] = find(parent[a])
    return parent[a]

# union 연산
def union(a, b):
    
    a = find(a) # a의 대표 노드 찾기
    b = find(b) # b의 대표 노드 찾기
    
    if a != b:
        parent[b] = a

# MST 실행
useEdge = 0  # 사용 엣지 수
result = 0  # 정답

while useEdge < (V - 1):

    v, s, e = heapq.heappop(pq)  # 큐에서 에지 정보 가져오기
    if find(s) != find(e):  # 사이클이 생기지 않으면
        union(s, e) # union 연산 수행
        result += v
        useEdge += 1

print(result)
 

 

[문제풀이]


이 문제는 MST 알고리즘을 코드로 구현할 수 있냐하는 문제다.

크루스칼 알고리즘을 이용, 에지를 가중치 순서대로 저장하고, 사이클이 생기지 않는 에지를 모두 연결하면 MST를 구할 수 있다.

 

 

  1. 에지 리스트로 그래프를 구현하고 유니온 파인드 리스트 초기화하기
  2. 그래프 데이터를 가중치 기준으로 정렬하기
  3. 가중치가 낮은 에지부터 연결 시도하기
  4. 연결한 노드의 개수가 N-1개가 될 때까지 3의 과정을 반복

 

# find 연산
def find(a):
    if a == parent[a]:
        return parent[a]
    else:
        parent[a] = find(parent[a]) # 경로 압축, 한번 find 연산을 할 때 중간경로도 다 대표노드로 바꾸게 된다. 
        return parent[a]
 

이때, 연결 시도할때 이렇게 find(a) 연산을 통해 사이클이 생기지 않는지 확인해야 한다.

두 노드의 대표노드가 같지 않으면 사이클이 생기지 않는다는 뜻이다.

그래서 find(a) 연산으로 유니온 파인드 리스트를 업데이트 시켜 두 노드의 대표 노드를 구한다.

 

# union 연산
def union(a, b):

    a = find(a) # a의 대표 노드 찾기
    b = find(b) # b의 대표 노드 찾기

    if(a != b): 
        parent[b] = a
 

find(a) 연산으로 두 노드의 대표 노드를 구했으면 그 후 union 연산으로 두 노드를 연결한다.

두 노드를 연결한다는 뜻은 b의 대표 노드를 a의 대표 노드로 설정한다는 뜻!

 

이 방식으로 모든 노드를 연결하는 최소 신장 트리를 구할 수 있다!