백준 16398 [행성 연결] 파이썬

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

 

# 백준 16398 행성 연결
# 문제 요약:
#   - N개의 행성을 모두 연결하는데 필요한 최소 비용을 구하는 문제
#   - 주어진 것은 N×N 행렬 (i→j 연결 비용)
#   - 최소 스패닝 트리(MST) 문제 (Kruskal / Prim 사용 가능)

import sys
sys.setrecursionlimit(10**7)
input = sys.stdin.readline

N = int(input())  # 행성(노드) 개수

# Union-Find parent 배열
parent = [i for i in range(N+1)]

# Find 함수 (경로 압축)
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

# Union 함수
def union(a, b):
    if a != b:
        parent[a] = b

# 모든 간선을 저장할 리스트
costs = []

# N×N 연결 비용 행렬 입력
for i in range(1, N+1):
    i_to_j_costs = list(map(int, input().split()))
    for j in range(1, N+1):
        if i == j:
            continue  # 자기 자신으로 가는 비용은 무시
        else:
            # (시작, 끝, 비용) 저장
            costs.append((i, j, i_to_j_costs[j-1]))

# 비용 기준 정렬 (오름차순 정렬해야 하므로 -x[2] 후 pop()으로 꺼냄)
costs.sort(key=lambda x: -x[2])

ans = 0  # 최종 최소 비용
# Kruskal 알고리즘: 작은 비용부터 간선 선택
while costs:
    a, b, cost = costs.pop()  # 가장 작은 비용 간선 꺼냄
    a = find(a)
    b = find(b)
    if a != b:                 # 같은 집합이 아니면 연결
        parent[a] = b
        ans += cost

print(ans)  # 최소 비용 출력

UNION-FIND를 이용하여 크루스칼 알고리즘을 구현하였다.

MST(최소 신장 트리) 구현 완료.

 

GPT가 알려준 프림 알고리즘을 이용한 풀이도 첨부한다.

visited와 최소힙을 이용하여 

방문하지않고, 현재 확인 되는 값들 중 가장 작은 cost의 값들을 뽑아 연결하는게 핵심이다.

# 백준 16398 행성 연결 - Prim 기반 최적화 버전
import sys
import heapq
input = sys.stdin.readline

N = int(input())

# 입력 행렬 저장
graph = [list(map(int, input().split())) for _ in range(N)]

visited = [False] * N   # MST에 포함된 노드 표시
min_heap = []           # 선택 가능한 간선 우선순위 큐

ans = 0
visited[0] = True  # 0번 노드부터 시작

# 0번 노드에서 연결 가능한 모든 간선을 힙에 삽입
for j in range(1, N):
    heapq.heappush(min_heap, (graph[0][j], j))

count = 1  # MST에 포함된 노드 수

while count < N:
    cost, node = heapq.heappop(min_heap)
    if visited[node]:
        continue
    visited[node] = True
    ans += cost
    count += 1
    # node에서 연결 가능한 간선 힙에 추가
    for j in range(N):
        if not visited[j]:
            heapq.heappush(min_heap, (graph[node][j], j))

print(ans)