[백준 2887 파이썬] 행성 터널 (플래티넘 5, MST)

알고리즘 유형 : MST
풀이 참고 없이 스스로 풀었나요? : X

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

소스 코드(파이썬)

import sys
input = sys.stdin.readline

# 유니온 파인드
def find(x):
    if parent[x] < 0:
        return x
        
    parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    root_a = find(a)
    root_b = find(b)
    
    if root_a == root_b:
        return False
    
    if parent[root_a] < parent[root_b]:
        parent[root_b] = root_a
    elif parent[root_a] > parent[root_a]:
        parent[root_a] = root_b
    else:
        parent[root_a] = root_b
        parent[root_b] -= 1
    
    return True

# 세 축을 기준으로 각각 정렬한 좌표에 대해,
# 인접한 노드 사이 간선을 모두 구해서 edges에 넣기
def findEdge(cdns, edges, axis):
    coordinates.sort(key=lambda x: x[axis])
    
    for idx in range(len(coordinates)-1):
        node1 = coordinates[idx]
        node2 = coordinates[idx+1]
        cost = abs(node1[axis] - node2[axis])
        
        edges.append((cost, node1[0], node2[0]))
    
    return

N = int(input())
coordinates = []
edges = []
parent = [-1]*(N+1)

# 입력 순서대로 노드 번호를 부여해줌
for num in range(1, N+1):
    x, y, z = map(int, input().split())
    coordinates.append((num, x, y, z))

# 축 별로 정렬 후 인접노드 간선 모두 구하기
for axis in range(1, 4):
    findEdge(coordinates, edges, axis)

# 크루스칼
edges.sort()
res = 0
for w, n1, n2 in edges:
    if union(n1, n2):
        res += w

print(res)

풀이 요약

1. 모든 노드를 N-1개의 간선으로 최소 비용으로 연결하는 문제이므로, 딱 봐도 최소 스패닝 트리 문제이다.

   간선의 가중치를 문제에서 정의해줬으므로, 모든 노드들 사이의 간선의 가중치를 구할 수 있다.

   그런데 이 문제에서 N은 10^5까지 주어지므로, 모든 노드들 사이의 간선을 다 구하게되면 그것만 해도 O(N^2)가 소모되어 TLE가 된다.

   즉, MST 알고리즘에 사용할 간선의 개수를 최소화해야한다는 뜻이다.

2. 어떤 식으로 간선을 구해야할까.

   간선의 가중치의 정의를 살펴보면, 두 좌표에 대해 어느 한 축의 좌표값의 차의 절댓값 중에, 가장 최소인 값을 가중치로 정한다.

   그러니 일단 각각의 축에 대해 생각해보자.

   MST 알고리즘에서는 가장 적은 비용으로, 정해진 개수만큼의 간선으로 모든 노드를 연결해야한다.

   즉, 같은 비용에 대해 최대한 많은 노드를 연결할수록 좋을 것이다.

   한편, 한 축의 좌표값을 기준으로 오름차순으로 정렬해보자.

   문제의 테스트 케이스를 예로 들면, x축을 기준으로 정렬하면

   -1 -1 -5
   10 -4 -1
   11 -15 -15
   14 -5 -15
   19 -4 19

   그런데, 만약 두 번째 노드에서 네 번쨰 노드로 간선을 이으면 비용이 4가 소요된다. 그러나 10에서 11로, 11에서 14로 이으면 똑같은 비용 4로 무려 노드 3개를 서로 연결시킬 수 있다.

   즉, 무조건 오름차순 정렬했을 때 인접한 노드들 사이의 간선이 MST 알고리즘의 후보가 된다.

3. 이제는 각 축 별로 정렬하여 모든 인접 간선을 구한 뒤, 이를 활용하여 크루스칼 알고리즘을 수행해주면 답을 구할 수 있다.

배운 점, 어려웠던 점

• 각 축 별로 정렬하는 것 까지는 생각해봤는데, 정렬된 좌표에서 인접 노드 사이의 간선들이 크루스칼 MST 알고리즘의 후보가 된다는 발상은 진짜 생각지도 못했다. 침착하게 크루스칼 MST 알고리즘을 배운 내용 안에서 가장 당연한 상식부터 출발해서 차근차근 사고를 개진시키면 충분히 찾아낼 수 있을 부분인 것 같다. 노력하자!