[백준 6091번] 핑크 플로이드 (Python3)

import sys
input = sys.stdin.readline
from heapq import *

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

N = int(input())

hq = []
for i in range(1,N):
  data = [*map(int,input().split())]
  for j in range(i+1,N+1):
    heappush(hq,(data[j-1-i],i,j))
parent = [i for i in range(N+1)]
tree = [[] for i in range(N+1)]
for i in range(N-1):
  while 1:
    w,x,y = heappop(hq)
    if findparent(x)==findparent(y):
      continue
    parent[parent[x]] = parent[y]
    tree[x].append(y)
    tree[y].append(x)
    break

for i in range(1,N+1):
  print(len(tree[i]),*sorted(tree[i]))

크루스칼 알고리즘 (MST) 로 해결하였다.

풀이가 쉽게 떠오르는 문제는 아니었다. 처음에는 문제 이름부터 플로이드니 플로이드-워셜로 풀 수 있겠다고 생각했다. 만약 플로이드 알고리즘에서 DP[i][j] = DP[i][k]+DP[k][j]가 나오면, i-j는 직접 연결되지 않았다는 뜻이므로 이를 이용할 수 있다. 하지만 N의 크기가 최대 1024이기 때문에, O(N^3)인 플로이드는 무리였다.

 

따라서 다른 방법을 생각하던 와중에, "트리는 간선이 N-1개" 라는 것이 떠올랐다. 이를 이용하면 가장 작은 간선부터 해서 N-1개를 연결하면 그게 곧 트리를 이루는 간선이지 않나 하고 생각했다. 풀이를 떠올리고 보니 "뭐야 생각해보니까 그게 바로 크루스칼 알고리즘이잖아" 하게 되었다.

 

크루스칼 알고리즘으로 트리를 연결하고, x,y를 연결시, x의 그래프에 y를, y의 그래프에 x를 넣어주면 트리가 완성된다.

 

 

여담) 이 문제의 숏코딩 순위 1위가 되었다.