[백준/Python]16398: 행성 연결

문제

https://www.acmicpc.net/problem/16398
백준 문제집 - 0x1B강 - 최소 신장 트리

 

알고리즘: 최소 스패닝 트리

---

풀이

크루스칼이 익숙한 터라 크루스칼 → 프림 순으로 풀었다.
크루스칼로도 수월히 통과됐지만, 이 문제는 "모든 간선"이 주어지므로 프림 알고리즘을 사용하는게 더 효율적이다.

 

크루스칼 알고리즘 사용

간선이 직접적으로 주어지는게 아니라(u v w), 행렬 자체로 주어진다.
그리고 A → B = B → A 양방향 그래프이므로 주대각선(i) 전까지만 체크해주면 된다.
(단방향이라면 모두 저장해줘야함)

 

그 다음은 일반적인 크루스칼 풀이와 같다.

1. 비용 기준으로 간선 정렬
2. 정렬된 간선들을 탐색
   • 같은 집합이 아니라면 연결 후 비용 저장

 

전체 코드

# 1. 크루스칼 알고리즘 풀이
# 메모리: 102560KB / 시간: 948ms
from sys import stdin


input = stdin.readline

N = int(input())
graph = []
parent = list(range(N))

for i in range(N):
    line = list(map(int, input().split()))
    # 주대각선(자기 자신) 전까지만 체크
    for j in range(i):
        graph.append((i, j, line[j]))

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]


def union(a, b):
    a, b = find(a), find(b)

    if a != b:
        if parent[a] < parent[b]:
            parent[b] = a
        else:
            parent[a] = b
        return True
    return False


graph.sort(key=lambda x: x[2])
ret = 0

for u, v, w in graph:
    # 같은 집합이 아니라면 결과값에 비용 추가
    if union(u, v):
        ret += w

print(ret)

---

프림 알고리즘 사용 (heapq, 리스트)

프림을 사용할경우 모든 간선들을 저장해둬야한다.
현재 선택된 노드를 기준으로 인접 노드들을 탐색해야하기 때문이다.

 

이 외에는 역시 일반적인 프림 풀이와 동일하다.

(힙을 사용할것인지, 리스트를 사용할것인지는 개인 선택)

 

힙을 사용했을때의 과정은 아래와 같다.

먼저 최소값들을 저장할 리스트distance를 생성해주고, heap에 임의의 노드와 값(0)을 집어넣는다.

1. heap에서 (비용, 노드)를 꺼낸다.
2. 이미 방문한 노드라면 넘어간다.
3. 아니라면 방문처리 후 아래의 과정을 실행한다.
   • 전체 비용에 꺼낸 비용을 더해줌
   • 확인한 노드의 수를 +1 카운트
     • 카운팅 후 노드의 수가 N개일경우 break
   • 현재 노드와 연결된 인접 노드들을 탐색 (비용, 노드번호)
     • 이미 방문한 노드거나 현재까지의 최솟값이 비용보다 작거나 같다면 넘어감
     • 아니라면, heap에 (비용, 노드번호)형태로 추가
     • 해당 비용값으로 최솟값 갱신 => distance[노드번호] = 비용

 

전체 코드

# 2. 프림 알고리즘 풀이 (heapq 사용)
# 메모리: 74296KB / 시간: 392ms
from sys import stdin
from heapq import heappush, heappop


input = stdin.readline

N = int(input())
graph = [tuple(map(int, input().split())) for _ in range(N)]
distance = [float("inf")] * N  # 현재까지 만들어진 MST와 연결될 수 있는 최소 비용
visited = [False] * N

ret = cnt = 0
heap = [(0, 0)]

while heap:
    cost, node = heappop(heap)

    if visited[node]:  # 이미 방문한 노드라면 패스
        continue

    visited[node] = True
    ret += cost
    cnt += 1

    # 카운팅한 노드의 갯수가 N개라면 break
    if cnt >= N:
        break

    # 다음 노드의 번호, 비용
    for nxt, cost in enumerate(graph[node]):
        if visited[nxt] or distance[nxt] <= cost:  # 이미 방문한 노드거나 최솟값이 cost보다 작다면 넘어간다.
            continue
        heappush(heap, (cost, nxt))
        distance[nxt] = cost


print(ret)

 

다음은 똑같은 프림 알고리즘이지만, 리스트를 사용하는 방식이다.

기본적인 플로우는 힙 사용 버전과 같다.

다만 리스트를 사용할경우 매번 "현재 MST에서 연결할 수 있는 노드 중 최소비용인 노드"를 탐색해야한다.

 

1. 새로 연결할 노드 탐색
2. 현재 MST를 기준으로 연결비용이 최소인 노드 선택
3. 해당 노드를 MST에 추가
   • 해당 노드 방문처리, 전체비용에 연결 비용 추가
4. 추가한 노드의 인접 노드들 탐색
   • 이미 방문한 노드거나 현재까지의 최솟값이 비용보다 작거나 같다면 넘어감
   • 아니라면 최솟값 갱신

 

전체 코드

# 3. 프림 알고리즘 풀이 (리스트 사용)
# 메모리: 71096KB / 시간: 504ms
from sys import stdin


input = stdin.readline

N = int(input())
graph = [tuple(map(int, input().split())) for _ in range(N)]
distance = [float("inf")] * N  # MST에 포함된 노드들로부터 각 노드까지의 최소 연결비용
visited = [False] * N

# 0번째 노드의 값을 0으로 설정 후 MST 생성 시작
distance[0] = 0
ret = 0

# 모든 노드를 MST에 포함시킬때까지 반복
for _ in range(N):
    # 현재 MST에 연결할 수 있는 가장 작은 비용의 노드를 찾아야함.
    min_dis = float("inf")
    min_node = -1

    for i in range(N):
        if not visited[i] and distance[i] < min_dis:
            min_dis = distance[i]
            min_node = i
    
    # MST에 포함시킨 뒤 결과값에 비용 추가
    visited[min_node] = True
    ret += min_dis

    # 새로 추가된 노드로부터 다른 노드들까지의 최소 연결 비용 갱신
    for nxt in range(N):
        if visited[nxt] or distance[nxt] <= graph[min_node][nxt]:
            continue
        distance[nxt] = graph[min_node][nxt]


print(ret)

---

기본적인 최소 신장 트리 문제다. 연습하기 좋은듯.