[백준] 17132번: 두더지가 정보섬에 올라온 이유 - 파이썬

🔈 문제

두더지가 정보섬에 올라왔다!

두더지는 정보섬 지하에 여러 채의 자택을 소유하고 있다. 두더지는 집 사이를 오가는 걸 좋아한다. 정보섬에는 총 N개의 두더지 집이 있으며 이 집들은 N-1개의 길로 연결되어 있다. 임의의 집에서 또 다른 집으로 가는 경로는 항상 유일하게 하나만 존재한다. 즉, 두더지 집들은 트리 형태로 모두 연결되어 있다. 어떤 길을 지날 때, 두더지는 W만큼의 만족도를 얻는다.

어느 날, 아래의 그림과 같은 집을 가진 두더지는 집1에서 집4로 이동했다. 이때 거치게 되는 집은 (1 → 2 → 3 → 4)이다. 두더지는 한 번 이동할 때마다, 이동경로에 포함되는 만족도들 중에서 가장 최소인 만족도를 얻는다. 즉, (1 → 4)의 경우에는 만족도를 2만큼, (6 → 2)의 경우에는 만족도를 3만큼 얻게 된다.

 

두더지는 갑자기 모든 집의 쌍 (a, b)에 대한 이동을 끝냈을 때 얻을 수 있는 만족도의 총합이 얼마인지 궁금해졌다. (a, b)와 (b, a)는 같은 경로로 친다. 즉, 위의 그림의 경우 (1-2, 1-3, 1-4, 1-5, 1-6, 2-3, 2-4, 2-5, 2-6, 3-4, 3-5, 3-6, 4-5, 4-6, 5-6)를 모두 이동했을때 얻는 만족도의 총합을 말한다.

하지만 두더지는 계산하다가 지쳐 이내 잠에 들고 말았다! 착한 여러분들이 잠든 두더지를 대신해 계산해 주자.

 

📝입력

첫째 줄에 집의 수 N이 주어진다. (1 ≤ N ≤ 100,000)
이후 N-1개의 줄에 걸쳐 X, Y, W가 주어진다. 이는 집 X와 Y가 연결되어 있고, 이 길을 지날 때의 만족도가 W라는 뜻이다. (1 ≤ X,Y ≤ N, 1 ≤ W ≤ 200) 

 

📑출력

문제에 제시된 이동을 끝냈을 때, 두더지가 느끼는 만족도의 총합을 출력한다.

 

📚 문제 풀이

 

해당 문제는 유니온 파인드를 사용해서 해결할 수 있다. 

입력으로 주어지는 트리는 MST이고 각각의 간선은 만족도를 가진다. 

이동 경로에 포함된 만족도들 중에서 가장 최소인 만족도를 취한다는 조건이 있으니 $u$와 $v$를 잇는 간선의 만족도가 $w$일 때, 이 $w$가 해당 MST의 최대인 만족도라면 이 값은 $u$와 $v$를 이을때를 제외하고 사용되지 않는다. (중복된 값이 없다는 가정하에)

반대로 예를 들어보면, 집합 $a$와 집합 $b$를 잇는 간선 $c$가 해당 트리에 대해 최소인 만족도라면 해당 만족도는 집합 $a$의 갯수 X 집합 $b$의 갯수 만큼 사용된다. 

 

이 점을 이용하여 문제를 해결하면 된다.

 

내림차순으로 간선을 정렬하고 각각 노드마다 현재 연결되어 있는 노드의 갯수를 저장해준다. 그리고 union시 각각의 갯수를 곱하여 만족도의 총합을 얻을 수 있다. 

 

💻 코드

import sys
input = sys.stdin.readline
 
def find(parent:list[int], x:int) -> int:
    if parent[x] != x:
        parent[x] = find(parent, parent[x])
    return parent[x]
 
def union(parent:list[int], a:int, b:int, w:int, child: list[int]) -> None:
    if a < b:
        parent[b] = a
        child[a] += child[b]
        child[b] = child[a]
    else:
        parent[a] = b
        child[b] += child[a]
        child[a] = child[b]
 
def solve():
    n = int(input())
    edges = [list(map(int,input().split())) for _ in range(n-1)]
    edges.sort(key=lambda x:-x[2])
    parent = [i for i in range(n+1)]
    child = [1]*(n+1)
    ans = 0
 
    for x,y,w in edges:
        x = find(parent, x)
        y = find(parent, y)
        ans += (child[x]*child[y])*w
 
        union(parent, x, y, w, child)
    return ans
 
print(solve())

 

🍭 시행착오

 

유니온 파인드의 아이디어 찾기가 재미있었다. 

 

🔗 문제링크 두더지가 정보섬에 올라간 이유

17132번: 두더지가 정보섬에 올라온 이유