[파이썬] 백준 2887 : 행성터널 (플레5)

[파이썬] 백준 2887 : 행성터널 (플레5)

2887번: 행성 터널

---

문제

---

풀이

0. 방향성 생각

• 지문만 읽으면 MST 문제임을 바로 알 수 있다.
• 이전 골드 문제와는 다르게 입력 크기가 10만이라서 모든 간선 정보를 받는 경우에는 n(n-1)/2 = 5*10^11이다.
• 특정 차원을 선택해서 위치별로 정렬시키고, 인접한 node끼리만 간선 정보에 추가해준다. (증명은 질문 게시판에)

1. 입력

from collections import defaultdict as dd
import sys
input = lambda : sys.stdin.readline().rstrip()

n = int(input())
info = [[] for _ in range(3)]
for i in range(n):
    loc = tuple(map(int,input().split()))
    for j in range(3):
        info[j].append((loc[j],i))

inf = 1e10
edges = dd(lambda:inf)
for j in range(3):
    info[j].sort()
    for i in range(n-1):
        p1,idx1 = info[j][i]
        p2,idx2 = info[j][i+1]
        edges[(idx1,idx2)] = min(abs(p1-p2),edges[(idx1,idx2)])
edges = list(edges.items())
edges.sort(key=lambda x:-x[1])

• dd에다가 인접 노드를 연결하는 경우 최소값으로 갱신해준다.

2. 함수 정의

parent = list(range(n))

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

def union(x,y):
    parent[max(x,y)] = min(x,y)

• union에는 항상 루트 노드를 넣고 작은 노드를 부모로 삼는다.
• 안에서 find를 해서 해도 되지만, union을 하는 조건문이 바깥에 있으면 중복해서 찾기 때문에 비효율적이다.

3. 출력

count,answer = 0,0
while count < n-1:
    nodes,cost = edges.pop()
    a,b = nodes
    pa,pb = find(a),find(b)
    if pa != pb:
        union(pa,pb)
        answer += cost
        count += 1
print(answer)

• n-1개를 연결했으면 루프를 탈출해준다.
• edges에서 만든 정보들을 pop시켜서 find를 통해 부모를 찾고, 다른 집합에 속해있으면 연결해준다.

 

---

전체코드

from collections import defaultdict as dd
import sys
input = lambda : sys.stdin.readline().rstrip()

n = int(input())
info = [[] for _ in range(3)]
for i in range(n):
    loc = tuple(map(int,input().split()))
    for j in range(3):
        info[j].append((loc[j],i))

inf = 1e10
edges = dd(lambda:inf)
for j in range(3):
    info[j].sort()
    for i in range(n-1):
        p1,idx1 = info[j][i]
        p2,idx2 = info[j][i+1]
        edges[(idx1,idx2)] = min(abs(p1-p2),edges[(idx1,idx2)])
edges = list(edges.items())
edges.sort(key=lambda x:-x[1])

parent = list(range(n))

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

def union(x,y):
    parent[max(x,y)] = min(x,y)

count,answer = 0,0
while count < n-1:
    nodes,cost = edges.pop()
    a,b = nodes
    pa,pb = find(a),find(b)
    if pa != pb:
        union(pa,pb)
        answer += cost
        count += 1
print(answer)

 

코멘트

증명은 모르겠고, 삼각 부등식 생각 + 직관으로 가까운거만 넣어서 풀긴 했는데 아직 와닿지는 않는..

유니온 시킬 때, 루트노드 넣어야되는데 그냥 노드 넣어서 답이 안나와서 한참 해맸음. 웃긴건 테케랑 반례는 맞음 ㅠ