[백준 2887] 행성 터널: 크루스칼 (python, 파이썬)

1.문제

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

2887번: 행성 터널

2.풀이

O(N^2)으로 풀리지 않는다. 각 차원마다 정렬하고, 가까이 있는 행성들끼리의 거리를 heap에 넣어준다.

될까 싶지만 되는 꿀팁 문제인 것 같다.

 

3.구현

import sys
import heapq

sys.stdin = open("input.txt")

N = int(sys.stdin.readline())

lis=[]
for i in range(N):
  x,y,z=map(int,sys.stdin.readline().split())
  lis.append((x,y,z,i))
# print(lis)
parent=[i for i in range(N+2)]

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

def uni(x,y,parent):
  x=find(x,parent)
  y=find(y,parent)
  if x==y:
    return
  parent[y]=x
  return

edges=[]
for i in range(3):
  lis.sort(key=lambda x:x[i])
  for idx in range(len(lis[:-1])):
    w = abs(lis[idx][i]-lis[idx+1][i])
    heapq.heappush(edges, (w,lis[idx][3],lis[idx+1][3]))

answer=0
cnt=0
while edges:
  cur = heapq.heappop(edges)
  if find(cur[1], parent)!=find(cur[2],parent):
    uni(cur[1],cur[2],parent)
    answer+=cur[0]
    cnt+=1

print(answer)