백준 2887: 행성 터널 (Python)

 

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

2887번: 행성 터널

 

처음에는 1774번(우주신과의 교감) 문제처럼 비용에 따른 간선을 모두 정의해주었다.

하지만 문제 조건에서 n의 수가 굉장히 크기 때문에 메모리 초과가 발생한다.

 

import sys
input = sys.stdin.readline

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

def union(parent, a, b):
    a = find(parent, a)
    b = find(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

n = int(input())

parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

planet = [0] * (n+1)
for i in range(1, n+1):
    planet[i] = list(map(int, input().split()))

edges = []
for i in range(1, n+1):
    for j in range(1, n+1):
        x = abs(planet[i][0] - planet[j][0])
        y = abs(planet[i][1] - planet[j][1])
        z = abs(planet[i][2] - planet[j][2])
        edges.append((min(x, y, z), i, j))
edges.sort()

res = 0
for i in edges:
    c, a, b = i
    if find(parent, a) != find(parent, b):
        union(parent, a, b)
        res += c

print(res)

 

메모리 초과 문제를 해결하려면 고려하는 간선의 수를 줄여야 한다.

문제 조건에서 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이기 때문에

 

1. x값 기준 오름차순 정렬

2. 서로 인접한 행성의 비용만을 고려한다. [i] ~ [i+1]

      ㄴ [i] ~ [i+2]는 [i] ~ [i+1] + [i+1] ~ [i+2] 와 같기 때문에 고려하지 않아도 된다.

3. 즉, n개의 행성에 대해 n-1개의 간선만 고려하면 된다.

4. y, z에 대해서도 위와 같은 과정을 반복한다.

 

이러한 방법으로 간선의 수를 줄일 수 있다.

 

import sys
input = sys.stdin.readline

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

def union(parent, a, b):
    a = find(parent, a)
    b = find(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

n = int(input())

parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

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

edges = []
for i in range(1, 4):
    planet.sort(key=lambda x: x[i])  # x값, y값, z값 기준 오름차순 정렬
    for j in range(n-1):
        edges.append((abs(planet[j+1][i]-planet[j][i]), planet[j][0], planet[j+1][0])) # 간선 정보 추가
edges.sort() # 비용 순으로 정렬

res = 0
for i in edges:
    c, a, b = i
    if find(parent, a) != find(parent, b):
        union(parent, a, b)
        res += c

print(res)