[벡준] 2887번: 행성 터널 - 파이썬(python)

틀린 풀이

import heapq

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

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

N = int(input())
arr = []
q = []
parent = [i for i in range(N)]

for _ in range(N) :
  x,y,z = map(int, input().split())
  arr.append((x,y,z))

for i in range(N) :
  for j in range(i+1,N) :
    cost = min(abs(arr[i][0]-arr[j][0]),abs(arr[i][1]-arr[j][1]),abs(arr[i][2]-arr[j][2]))
    heapq.heappush(q,(cost,i,j))

answer = 0
while q :
  cost,i,j = heapq.heappop(q)
  if find_parent(parent,i) != find_parent(parent,j) :
    union_parent(parent, i, j)
    answer += cost

print(answer)

간선을 구하기 위해 2중 for문을 사용하여 풀었다.
1번 노드와 연결된 2,3,...n-1 노드와의 간선을 모두 구한 후
2번 노드와 연결된 3,4,...n-1 노드와의 간선을 구하는..
즉, N^2 시간에 모든 간선을 구한 방법이였다.
그러나 이 방법은 메모리 초과로 실패하게 되었다.
행성의 개수가 최대 100000개인데, 이 방법은 n(n+1)/2 로 약 5000000000 의 시간이 걸릴 수 있기 때문이다!
따라서 간선을 구하기 위한 다른 방법을 생각해야 한다.

맞은 풀이

import heapq
import sys
input = sys.stdin.readline

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

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

N = int(input())
x_arr = []
y_arr = []
z_arr = []
q = []
parent = [i for i in range(N)]

for i in range(N) :
  x,y,z = map(int, input().split())
  x_arr.append((x,i))
  y_arr.append((y,i))
  z_arr.append((z,i))

x_arr.sort()
y_arr.sort()
z_arr.sort()

for i in range(N-1) :
  heapq.heappush(q,(x_arr[i+1][0] - x_arr[i][0], x_arr[i][1], x_arr[i+1][1]))
  heapq.heappush(q,(y_arr[i+1][0] - y_arr[i][0], y_arr[i][1], y_arr[i+1][1]))
  heapq.heappush(q,(z_arr[i+1][0] - z_arr[i][0], z_arr[i][1], z_arr[i+1][1]))

answer = 0
while q :
  cost,i,j = heapq.heappop(q)
  if find_parent(parent,i) != find_parent(parent,j) :
    union_parent(parent, i, j)
    answer += cost

print(answer)

터널의 비용이 min(|xA-xB|, |yA-yB|, |zA-zB|) 인 것을 이용한 풀이이다.
즉, x축을 고려해서 정렬한 것/ y축을 고려해서 정렬한 것 / z축을 고려해서 정렬한 것 에 대해 각각 수행하고,
오름차순으로 정렬된 ?축에 대해 (i+1번째 ?축 위치 - i번째 ?축 위치) 값을 간선으로 저장한다.
항상 i번째 위치와 가장 근접한 i+1번째 위치와의 차이가 가장 작은 거리가 될 것이므로 이 식이 성립한다.

만약, 1행성과 2행성간의 거리가 x축에 의해서도 저장되고, y혹은 z축에 의해서도 저장된다면 중복되지 않을까?

• 이미 x,y,z간의 거리 중 더 작은 비용에 의해 union_parent() 연산이 수행되었다면, 1행성과 2행성은 같은 부모를 가지기 때문에 뒤에 나오는 중복된 간선에 대해서는 처리를 하지 않을 것이다!

이 방법은 각각의 축에 대해 N의 시간이 소요되므로 3N 만에 간선을 구할 수 있는 풀이이다!
N^2 -> 3N으로 시간이 엄청 줄었다!