[백준] 2887번 행성 터널

[나중에 다시 풀기]

그냥 하던대로 거리 구해서 모든 간선을 추가하면 메모리 초과가 난다.

<메모리 초과 코드>

from itertools import combinations
import sys
input = sys.stdin.readline

n = int(input())
xyz_ = [[0,0,0]]
for _ in range(n):
    xyz_.append(list(map(int, input().split())))

pair_ = list(combinations([i for i in range(1, n+1)], 2))
edges = []
answer = 0
parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

def distance(a, b):
    return min(abs(a[0]-b[0]), abs(a[1]-b[1]), abs(a[2]-b[2]))

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

def union_node(a, b):
    a = find_parent(a)
    b = find_parent(b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

for i in pair_:
    edges.append([i[0], i[1], distance(xyz_[i[0]], xyz_[i[1]])])

edges = sorted(edges, key = lambda x: x[2])

for edge in edges:
    if find_parent(edge[0]) != find_parent(edge[1]):
        union_node(edge[0], edge[1])
        answer = answer + edge[2]

print(answer)

모든 간선이 아니라, x, y, z 에 대해 각각 거리를 구해서 최소거리 순으로 정렬한 결과를 사용하면 된다.

ex) N개의 점, 기존 방식대로라면  N(N+1)/2개의 간선에 대한 정보를 전부 계산하고 저장해야함.

그러나, x, y, z를 따로 구하면 3N만큼의 정보만 저장해도 됨. 

 

<통과 코드>

from itertools import combinations
import sys
input = sys.stdin.readline

n = int(input())
x_ = []
y_ = []
z_ = []
parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

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

def union_node(a, b):
    a = find_parent(a)
    b = find_parent(b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

for i in range(1, n+1):
    x, y, z = map(int, input().split())
    x_.append([x, i])
    y_.append([y, i])
    z_.append([z, i])

x_ = sorted(x_)
y_ = sorted(y_)
z_ = sorted(z_)
edges = []
for i in range(n-1):
    edges.append([abs(x_[i][0]-x_[i+1][0]), x_[i][1], x_[i+1][1]])
    edges.append([abs(y_[i][0]-y_[i+1][0]), y_[i][1], y_[i+1][1]])
    edges.append([abs(z_[i][0]-z_[i+1][0]), z_[i][1], z_[i+1][1]])

edges = sorted(edges)
answer = 0
for edge in edges:
    cost, p1, p2 = edge
    if find_parent(p1) != find_parent(p2):
        union_node(p1, p2)
        answer = answer + cost
print(answer)

이런 생각은 도대체 누가 하는 거지... 

x, y, z 별로 따로 거리를 저장해서, 그 거리를 sorting한 다음, 제일 작은 것 부터 꺼낸다라... 아주 ... 아주 훌륭한 생각이다.