[Baekjoon] 2887번 : 행성 터널 (Python, 최소 신장 트리)

[Baekjoon] 2887번 : 행성 터널 (Python, 최소 신장 트리)

 

목차

 

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

2887번: 행성 터널

코드 링크

https://github.com/stellaluminary/Baekjoon

GitHub - stellaluminary/Baekjoon

Python

 

(실패) 방법 1 - 메모리 초과

 

일반적인 최소 신장 트리라고 판단할 수 있다.

N개의 행성(노드)가 있고 터널(간선)이 존재하며 비용(cost)을 최소화하는 문제이다.

이에 대해 크루스칼 알고리즘을 활용하여 문제를 풀면 edges마다의 cost를 구하기 위해 $O(N^2)$를 구하기 때문에  메모리 문제가 발생한다.

 

import sys
input = sys.stdin.readline

def f_p(p,x):
    if p[x] != x:
        p[x] = f_p(p, p[x])
    return p[x]

def union(p,a,b):
    a, b = f_p(p,a), f_p(p,b)
    p[max(a,b)] = min(a,b)

def cost(c1, c2):
    return min(abs(c1[0]-c2[0]), abs(c1[1]-c2[1]), abs(c1[2]-c2[2]))

n = int(input())
p = [i for i in range(n+1)]
coord=[]
edges=[]
res=0

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

for i in range(n-1):
    for j in range(i+1, n):
        edges.append((cost(coord[i], coord[j]), i+1, j+1))

edges.sort()

for edge in edges:
    c, a, b = edge
    if f_p(p, a) != f_p(p, b):
        union(p, a, b)
        res += c

print(res)

 

방법 2 - 메모리 88332KB / 시간 1416ms / 코드 길이 649B

 

결국 계산하는 간산의 숫자를 줄여야 한다.

기존에는 ${V(V-1) \over 2}$개의 간선을 메모리에 저장하지만 잘 생각해보면 꼭 우리가 모든 간선의 cost를 구할 필요가 없다.

우리가 N개의 노드가 존재한다면 결국 N-1개의 간선 연결만 있으면 모든 통로가 연결된 것이나 마찬가지 이기 때문이다.

그러나 여기서 주의해야할 점은 그럼 N-1개만 뽑으면 되냐하면 그렇지 않다.

본 문제에서는 cost를 구하기 위해 3개의 좌표가 주어지고 각 축에 대한 별도의 계산을 진행하기 때문에 축별로 N-1 후보군을 각각 뽑아야 한다.

그 이유는 한 축에 대한 cost가 다른 축 대비 다 작을 수 있기 때문이다.

예를 들어 x축 좌표에 대한 N-1 간선들의 비용이 모두 0 이고 y축, z 축에 대한 비용이 0보다 크다면 당연히 x 축에 대한 간선들만 뽑기 때문이다.

고로 우리가 해야하는 일은 각 축에 대한 정렬로 후보군을 N-1개씩 뽑아 크루스칼 알고리즘을 돌리면 된다.

 

 

import sys
input = sys.stdin.readline

def f_p(p,x):
    if p[x] != x:
        p[x] = f_p(p, p[x])
    return p[x]

def union(p,a,b):
    a, b = f_p(p,a), f_p(p,b)
    p[max(a,b)] = min(a,b)

n = int(input())
p = [i for i in range(n+1)]
coord=[]
edges=[]
res=0

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

for i in range(3):
    coord.sort(key=lambda x:x[i])
    for j in range(1, n):
        edges.append((abs(coord[j-1][i] - coord[j][i]), coord[j-1][3], coord[j][3]))

edges.sort()

for edge in edges:
    c, a, b = edge
    if f_p(p, a) != f_p(p, b):
        union(p, a, b)
        res += c

print(res)