[BOJ] 16398. 행성 연결

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

16398번: 행성 연결

 

 

1. 해결방법

"""
1. 아이디어
- MST 이용
- 행성의 개수를 입력, 행성 간의 거리 비용를 2차원 배열로 입력 (Cij)
- [비용, 목적지 행성]의 형태로 행성의 인접리스트를 생성
- heap이 빌때까지 다음을 반복
    - 힙에서 최솟값을 꺼내고 방문 여부를 체크
    - 방문하지 않았다면 result에 비용을 추가, 방문을 했다면 pass
    - 인접리스트에 인덱스 값을 가져와서 방문하지 않은 행성은 heap에 push

2. 시간복잡도
- O(nlogn)

3. 자료구조
- heap : [비용, 목적지 행성]
- 인접리스트 : [비용, 목적지 행성]
- 방문 여부 : chk[]
- 비용 합 : int
"""

 

 

2. 정답코드

import sys
import heapq
input = sys.stdin.readline

N = int(input())

# 인접리스트
edge = [[]]

# 행성 거리 input
for _ in range(N):
    edge.append(list(map(int, input().split())))

heap = [[0, 1]]
result, count = 0, 0
chk = [False] * (N + 1)

while heap:
    if count == N:
        break
    w, each_node = heapq.heappop(heap)

    if chk[each_node] == False:
        chk[each_node] = True
        result += w
        count += 1

        for i in range(N):
            if each_node - 1 == i:
                continue
            heapq.heappush(heap, [edge[each_node][i], i + 1])

print(result)

MST 기본 문제에서 약간 변형한 문제이다.
개인적으로 시간이 오래 걸렸으며, 시간초과로 인해서 수정을 거듭한 코드이다.