[COCI] 행성 터널

출처: https://www.acmicpc.net/problem/2887

2887번: 행성 터널

---

 무향그래프에 대해서 최소 신장트리를 구하는 문제이므로, 크루스칼 알고리즘을 적용하면 쉽게 풀 수 있습니다.

크루스칼 알고리즘: https://jcy1996.tistory.com/35

크루스칼 알고리즘(Kruskal Algorithm)

 

다음은 소스코드입니다.

from itertools import combinations

n = int(input())
array = []
for i in range(n):
	array.append([i] + list(map(int, input().split())))

graph = []
for a, b in combinations(array, 2):
    i1, x1, y1, z1 = a
    i2, x2, y2, z2 = b
    graph.append((min(abs(x1 - x2), abs(y1 - y2), abs(z1 - z2)), i1, i2))
    
graph.sort()

parent = [i for i in range(n)]

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

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        
result = 0
for d, x, y in graph:
    if find_parent(parent, x) != find_parent(parent, y):
        union_parent(parent, x, y)
        result += d
        
print(result)

 위는 가장 기본적인 크루스칼 알고리즘을 구현한 것입니다. 다만, 다음과 같이 구현할 경우 가능한 모든 조합에 대해 순회하게 되므로, 시간복잡도는 $O(N^{2})$이 됩니다. 행성의 개수가 10만 개 이하이므로 최악의 경우 100억 번을 연산하게 될 수도 있다는 말입니다. 따라서 좀 더 효율적인 코드 작성이 필요합니다.

 아래는 개선된 코드입니다.

n = int(input())
array = []
for i in range(n):
	array.append([i] + list(map(int, input().split())))

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

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

parent = [i for i in range(n)]
graph = []

for i in range(1, 4):
    tmp = sorted([(x[i], x[0]) for x in array])
    graph += [(tmp[j + 1][0] - tmp[j][0], tmp[j][1], tmp[j + 1][1]) for j in range(n - 1)]
    
graph.sort()

for c, x, y in graph:
    if find_parent(parent, x) != find_parent(parent, y):
        union_parent(parent, x, y)
        result += c
print(result)

 사실 필드가 좌표축인 시점에서 우리는 모든 행성을 연결하는 경우의 수에 대해 고려할 필요가 없습니다. 좌표축은 선형적이기 때문입니다. 즉, 행성이 x축 기준 -1, 3, 10, 5와 같이 놓여있다고 가정했을 경우 6가지의 모든 조합을 고려하는 것이 아니라  (-1에서 3으로 연결, 3에서 5로 연결, 5에서 10으로 연결)과 같이 3가지의 간선만 고려하면 됩니다.

 따라서 각각의 축을 기준으로 최소비용순으로 정렬하여 간선 정보를 초기화하기로 하였습니다.

for i in range(1, 4):
    tmp = sorted([(x[i], x[0]) for x in array])
    graph += [(tmp[j + 1][0] - tmp[j][0], tmp[j][1], tmp[j + 1][1]) for j in range(n - 1)]
    
graph.sort()

 마지막 순회에서는 모든 축에 대해 가장 비용이 짧은 간선부터 추가하므로, 최소 신장트리를 보장받을 수 있습니다.