[플래티넘 5] 백준 20188 - 등산 마니아 (파이썬)

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

풀이

트리 DP 또는 서브트리 크기 기반의 조합 계산을 사용한 문제이다.

• 무방향 트리이며, 등상 마니아들은 경로 (a, b, c)를 선택할 때 항상 a → b → c 순으로 이동한다.
• 단, b는 a와 c의 LCA(최소 공통 조상)이어야 한다.
• 모든 (a, b, c) 쌍에서의 경우의 수를 계산해야한다.

사용한 아이디어

1. 어떤 간선을 제거하면 트리가 두 개의 서브트리로 나뉘게 된다.
2. 해당 간선을 통해 만들어지는 (a, b, c)의 경우의 수는 다음 두개로 나뉜다.
   1. 서브 트리 내부에서 a, c를 고르고 b가 그 루트인 경우 → sizes[i] * (sizes[i] - 1) / 2
   2. 서브트리와 바깥에서 a, c를 하나씩 고르고 b가 그 경계인 경우 → sizes[i] * (N - sizes[i])ㅊ

visit = [False] * N
sizes = [0] * N

def dfs(u) :
    visit[u] = True
    sizes[u] = 1
    for v in adj[u] :
        if not visit[v] :
            dfs(v)
            sizes[u] += sizes[v]

• 각 노드의 서브트리 크기를 구한다. (sizes[u])
• DFS로 트리를 탐색하면서 리프부터 크기를 누적한다.

answer = 0
for i in range(1, N) :
    answer += sizes[i] * (sizes[i] - 1) // 2        # 같은 서브트리 내에서 (a, c) 쌍
    answer += sizes[i] * (N - sizes[i])             # 다른 트리 간 (a, c) 쌍

• i는 부모 노드 0에서ㅓ 시작한 DFS를 기준으로, i는 항상 0의 하위에 위치한 노드이다.
• sizes[i]는 i가 루트인 서브트리의 크기이다.
• 이 서브트리의 노드 중 서로 다른 두 개의 노드를 (a, c)로 고르는 경우의 수는 다음과 같다.

결국. 이 코드는 트리에서 가능한 (a, b, c) 조합 수를 간선 단위로 세는 아이디어를 사용한 고급 트리 DP 문제 풀이다.

각 간선을 기준으로 서브트리를 나누어 가능한 조합 수를 더해가는 방식으로 최적화되어 있다.

코드

# 백준 20188 - 등산 마니아
# 분류 : 다이나믹 프로그래밍

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 6)

N = int(input())
adj = [[] for _ in range(N)]

for _ in range(N - 1) :
    a, b = map(int, input().split())
    a -= 1
    b -= 1

    adj[a].append(b)
    adj[b].append(a)

visit = [False] * N
sizes = [0] * N

def dfs(u) :
    visit[u] = True
    sizes[u] = 1

    for v in adj[u] :
        if not visit[v] :
            dfs(v)
            sizes[u] += sizes[v]

dfs(0)

answer = 0
for i in range(1, N) :
    answer += sizes[i] * (sizes[i] - 1) // 2
    answer += sizes[i] * (N - sizes[i])

print(answer)