백준(BOJ) 2887 행성 터널(Python)

문제 내용

 

최소 신장 트리 문제.

입력받은 행성을 각각 x, y, z좌표 기준으로 오름차순 정렬한 리스트를 3개 만들고, 각 리스트 별로 i번째 행성와 i+1번째 행성을 잇는 터널의 연결 비용을 계산해 전체 간선 리스트에 저장했다.

그 후 전체 간선 리스트에서 크루스칼을 돌렸다.

좌표 별로 정렬한 이유는, 결국 행성과 행성 간의 터널의 연결 비용은 행성의 세 좌표 중 하나에만 의존하기 때문이다.

행성 간의 모든 간선을 연결하면 당연히 시간 초과가 터질 것이기에, 간선의 수를 위와 같이 줄이면 총 3*(n-1)개의 간선이 만들어진다.

위와 같이 간선을 만들어도 상관 없는 이유는, 예를 들어 x1, x2, x3 순으로 리스트에 들어있을 때 (x1과 x2를 잇는 간선의 가중치)+(x2과 x3를 잇는 간선의 가중치)=(x1과 x3를 잇는 간선의 가중치)가 될 것이다. 그러나 좌변의 경우에는 간선을 2개 만들 수 있지만, 우변의 경우에는 간선을 1개밖에 만들지 못하므로 효율이 떨어지는 간선이라고 할 수 있다. 따라서 우변의 경우는 고려하지 않아도 되기 때문에 위와 같이 간선을 만들어도 상관이 없는 것이다.

 

코드는 다음과 같다.

 

# -*- coding: utf-8 -*-
import sys
from collections import deque
import heapq
import copy
#input=sys.stdin.readline
#sys.setrecursionlimit(100000000)

n=int(input())
L=[]
L1=[]
L2=[]
L3=[]
cycle=[i for i in range(n+1)]
ans=0
k=0
for i in range(n):
    x=list(map(int,input().split()))+[i+1]
    L1.append(x)
    L2.append(x)
    L3.append(x)
L1.sort()
L2.sort(key=lambda x:x[1])
L3.sort(key=lambda x:x[2])
for i in range(n-1):
    heapq.heappush(L,(abs(L1[i][0]-L1[i+1][0]),L1[i][3],L1[i+1][3]))
    heapq.heappush(L,(abs(L2[i][1]-L2[i+1][1]),L2[i][3],L2[i+1][3]))
    heapq.heappush(L,(abs(L3[i][2]-L3[i+1][2]),L3[i][3],L3[i+1][3]))
    
def find(u):
    if u!=cycle[u]:
        cycle[u]=find(cycle[u])
    return cycle[u]

def union(u,v):
    u=find(u)
    v=find(v)
    cycle[v]=cycle[u]

while k<n-1:
    c,a,b=heapq.heappop(L)
    if find(a)!=find(b):
        ans+=c
        union(a,b)
        k+=1
print(ans)

 

그래도 간선의 수가 많았는지 python으로 제출하면 시간 초과가 터졌다.

pypy로 제출해서 통과했다.