즐거운 화요일. MST와 함께 시작해보자.

BOJ 6091 - 핑크 플로이드 링크
(2022.10.04 기준 G1)
(치팅하면 영원히 금요일 안옴)

문제

1~N까지의 번호가 매겨져 있는 N개의 정점과 N-1개의 간선으로 이루어진 트리가 있고, 두 정점간 최단거리가 저장되어 있는 인접 행렬이 있을 때
인접 리스트 형태로 원래 트리를 표현

알고리즘

N-1개의 간선을 가지고 있는 트리는 스패닝 트리이며, 최단거리로 인접 행렬이 채워져 있으므로 이를 이용하여 MST를 만들면 된다.

풀이

플로이드 워셜을 통해 두 정점간의 최단거리가 인접 행렬로 채워져 있고, N개의 정점과 N-1개의 간선으로 이루어져 있던 트리이므로, 최단거리를 이용하여 N-1개의 간선을 가진 트리를 만드는 최소 스패닝 트리 문제임을 알 수 있다.

난 크루스칼 알고리즘을 통해 풀었으니 크루스칼로 설명하겠다.

먼저 인접행렬 형태로 주어지는 것을 입력받아야 한다.
N-1개의 줄로부터 i번 줄에는 i+1, i+2, ..., N 와의 최단 거리가 주어지므로
for문을 1부터 N-1까지, 주어지는 리스트는 i+1부터 N까지로 묶어서 입력받으면 된다.

edges = [] # 간선
for i in range(1, N):
	for j, 거리 in zip(range(i + 1, N + 1), list(map(int, input().split()))):
    	edges.append((i, j, 거리))

이런 느낌으로 edges에 저장하자.

그 다음은 간단하다. 여느 크루스칼 문제처럼 edges를 가중치가 오름차순이 되게끔 정렬하고, union-find 함수를 이용하여 mst를 만들면 된다. 만들어갈 때에는 결과를 저장할 인접 리스트에 꼭 추가하자.

코드

import sys; input = sys.stdin.readline

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

def union(u, v):
    u = find(u)
    v = find(v)
    if u < v:
        parent[v] = u
    else:
        parent[u] = v

N = int(input())
edges = [] # 간선
# i번 줄에는 i + 1, i + 2, ..., N번과의 최단거리가 주어진다.
for u in range(1, N):
    for v, w in zip(range(u + 1, N + 1), map(int, input().split())):
        edges.append([u, v, w])

parent = list(range(N + 1)) # 부모 노드
ct = 0 # 연결된 간선 수
result = [[] for _ in range(N + 1)] # 인접 리스트 형태로 결과 저장
for u, v, w in sorted(edges, key = lambda x: x[2]): # 가중치가 오름차순이 되게 정렬
    if find(u) != find(v): # 부모 노드가 다르면 union
        union(u, v)
        result[u].append(v) # 인접 리스트에 결과 저장
        result[v].append(u)
        ct += 1
        if ct == N - 1: # 연결된 간선 수가 N - 1개면 MST 완성이므로 중지
            break

# 1번부터 결과 출력
for i in range(1, N + 1):
    print(len(result[i]), *sorted(result[i]))

여담

개인적으론 이 문제가 왜 G1인지 모르겠다.
간선의 개수가 N - 1개라고 했으니 스패닝 트리이고, u -> v로 직접 이어진다면, 그 경로는 u와 v의 최단거리고 또 유일해야 하므로, 결국은 두 정점 사이의 간선을 유일하게끔 뽑아가면서 트리를 완성하는 크루스칼 돌리면 되는 것이다.
원래는 P5였던데.. 흠..