[Python][백준 2887][최소스패닝트리][Union-Find] 행성 터널 - 컴도리돌이

반응형

2887번: 행성 터널

---

풀이 과정

1. 최소 스패닝 문제로 풀었기에 입력받은 x, y, z와 인덱스 번호를 붙여서 before 배열에 삽입해주었다.
2. before 배열을 이중 반복문으로 다음과 같이 after 배열을 생성해주었다.
    2-1. x, y, z 순으로 정렬하기 위해서 첫 번째 반복문의 범위를 3만큼 돌려서 각각 x, y, z 순으로 정렬해주었다.
    2-2 두 번째 반복문은 범위를 1~n까지 설정하고 0번째 인덱스 값을 before_loc에 저장한다.
    2-3 before_loc과 after_loc  그리고 두 인덱스의 거리 차 절대 값을 after 배열에 저장한다. -> insert([두 인덱스 거리, 첫 번째 인덱스, 두 번째 인덱스])
    2-4 before_loc을 after_loc으로 저장한다.
3. x, y, z 순으로 정렬해서 거리, 출발 인덱스, 도착 인덱스 가 저장되어 있는 after 배열을 거리가 짧은 순으로 정렬시킨다.
4. after 배열의 첫 번째 인덱스부터 반복문을 진행해서 출발 인덱스와 도착 인덱스가 연결이 되어 있지 않으면 거리를 ans에 더해주고, cnt를 1만큼 더해준다.
5. n-1만큼 연결시키면 되기 때문에 cnt가 n-1만큼 연결되면 종료된다. (최소 스패닝 알고리즘이 어떻게 작동하는지는 생략.)

---

풀이 코드

import sys; input = sys.stdin.readline
def findParent(x) :
  if root[x] =! x : return findParent(root[x])
  return root[x]
def union(x,y) :
  x,y = findParent(x),findParent(y)
  if x != y :
    root[y] = x
    return True
  else : return False
n = int(input())
before,after = [list(map(int,input().split())) +[i] for i in range(n)],[]
root = [i for i in range(n)]
ans,cnt = 0,0
for i in range(3) :
  before.sort(key = lambda x : x[i])
  before_loc = before[0][-1]
  for j in range(1,n) :
    after_loc = before[j][-1]
    after.append([abs(before[j][i] - before[j-1][i]),before_loc,after_loc])
    before_loc = after_loc
for dist,start,end in sorted(after,key = lambda x : x[0]) :
  if union(start,end) : ans += dist ;cnt += 1
  if cnt == n-1:break
print(ans)