[백준 2887번 ] 행성 터널 문제 풀이 (with Python)

< 문제 >

 

 

< 입출력 >

 

 

입출력은 동일하나, 메모리 초과가 뜬다. N이 100,000인데 128MB이므로 info의 메모리를 줄여야한다.

< 메모리 초과 코드 >

import sys
input = sys.stdin.readline

def find(a):      # 루트 노드 찾기
    if a == parent[a]:
        return a
    parent[a] = find(parent[a])
    return parent[a]


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



N = int(input())
arr = []
info = []
result = 0
parent = [i for i in range(N)]
for _ in range(N):
    x, y, z = map(int, input().split())
    arr.append((x, y, z))

for i in range(N):
    for j in range(i+1, N):
        info.append((i, j, min(abs(arr[i][0]-arr[j][0]), abs(arr[i][1] -arr[j][1]), abs(arr[i][2] - arr[j][2]))))

info.sort(key=lambda x:x[2])

for a, b, c in info:
    if find(a) != find(b):
        union(a, b)
        result += c

print(result)

 

 

< 수정한 코드  >

import sys
input = sys.stdin.readline

def find(a):      # 루트 노드 찾기
    if a == parent[a]:
        return a
    parent[a] = find(parent[a])
    return parent[a]


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



N = int(input())

arr = []
info = []
result = 0
parent = [i for i in range(N)]
for i in range(N):
    x, y, z = map(int, input().split())
    arr.append((x, y, z, i))           # 인덱스를 같이 info 배열에 넣어줌

for i in range(3):
    arr.sort(key=lambda x: x[i])
    for j in range(1, N):
        info.append((abs(arr[j - 1][i] - arr[j][i]), arr[j - 1][3], arr[j][3]))

info.sort()


for a, b, c in info:
    if find(b) != find(c):
        union(b, c)
        result += a

print(result)