[백준] 16398번 행성 연결 (Python)

16398번: 행성 연결

문제

홍익 제국의 중심은 행성 T이다. 제국의 황제 윤석이는 행성 T에서 제국을 효과적으로 통치하기 위해서, N개의 행성 간에 플로우를 설치하려고 한다.

두 행성 간에 플로우를 설치하면 제국의 함선과 무역선들은 한 행성에서 다른 행성으로 무시할 수 있을 만큼 짧은 시간만에 이동할 수 있다. 하지만, 치안을 유지하기 위해서 플로우 내에 제국군을 주둔시켜야 한다.

모든 행성 간에 플로우를 설치하고 플로우 내에 제국군을 주둔하면, 제국의 제정이 악화되기 때문에 황제 윤석이는 제국의 모든 행성을 연결하면서 플로우 관리 비용을 최소한으로 하려 한다.

N개의 행성은 정수 1,…,N으로 표시하고, 행성 i와 행성 j사이의 플로우 관리비용은 Cij이며, i = j인 경우 항상 0이다.

제국의 참모인 당신은 제국의 황제 윤석이를 도와 제국 내 모든 행성을 연결하고, 그 유지비용을 최소화하자. 이때 플로우의 설치비용은 무시하기로 한다.

입력

입력으로 첫 줄에 행성의 수 N (1 ≤ N ≤ 1000)이 주어진다.

두 번째 줄부터 N+1줄까지 각 행성간의 플로우 관리 비용이 N x N 행렬 (Cij), (1 ≤ i, j ≤ N, 1 ≤ Cij ≤ 100,000,000, Cij = Cji, Cii = 0) 로 주어진다.

출력

모든 행성을 연결했을 때, 최소 플로우의 관리비용을 출력한다.

 

과정

1. Minimum Spanning Tree(MST)는 문제 프림알고리즘을 이용하여 풀이한다.
2. 방문을 확인하고 result에 비용을 더해준 후 출력한다.

 

예제 입력

3
0 2 3
2 0 1
3 1 0

예제 출력

3

 

풀이

import heapq
import sys
input = sys.stdin.readline

n = int(input())
flow = [list(map(int, input().split())) for _ in range(n)]
graph = [[] for _ in range(n)]
for i in range(n):
    for j in range(i+1, n):
        graph[i].append((flow[i][j], j))
        graph[j].append((flow[i][j], i))

visited = [False] * n

q = []
for i in graph[0]:
    heapq.heappush(q, i)

result = 0

while q:
    cost, node = heapq.heappop(q)

    if not visited[node] and node > 0:
        result += cost
        visited[node] = True
        for i in graph[node]:
            if not visited[i[1]]:
                heapq.heappush(q, i)

print(result)

 

한 줄씩 해석해보기

n = int(input())
flow = [list(map(int, input().split())) for _ in range(n)]
graph = [[] for _ in range(n)]
for i in range(n):
    for j in range(i+1, n):
        graph[i].append((flow[i][j], j))
        graph[j].append((flow[i][j], i))

• 행성의 수 n을 입력받는다.
• 플로우를 입력 받고 graph에 n x n 행렬을 만들어준다.

visited = [False] * n

q = []
for i in graph[0]:
    heapq.heappush(q, i)

result = 0

• 방문을 체크하기 위한 visited를 False로 만들어준다.
• graph의 0번째 값을 가져와 heap화 해준다
• 결과를 출력할 result를 만든다.

while q:
    cost, node = heapq.heappop(q)

• q에 값이 없을때까지 하나씩 가져와서 확인한다.
• 각각 cost와 node변수에 넣어준다.

    if not visited[node] and node > 0:
        result += cost
        visited[node] = True
        for i in graph[node]:
            if not visited[i[1]]:
                heapq.heappush(q, i)

print(result)

• 만약 visited의 값이 False이고 node가 0보다 크다면 True로 바꿔준다.
• result에 비용을 추가한다.
• 그 다음 노드를 넣어주고 방문하지 않은 노드를 계속 확인한다.