https://jungol.co.kr/problem/4602
| 시간 제한 | 메모리 제한 |
| 2초 | 128 MB |
문제
동네 뒷 산에는 등산로가 있다. 등산로는 $N$개의 작은 오두막들이 $N−1$개의 오솔길로 이어진 형태이다.
한 오솔길은 두 개의 오두막을 양 방향으로 연결한다. 한 오솔길의 길이는 1이다.
어떤 오두막에서도 다른 모든 오두막으로 하나 이상의 오솔길을 따라 이동하는 것이 가능하다.
오두막들은 1번부터 $N$번까지 번호가 붙어있으며, 1번 오두막이 산 정상에 있다.
1번 오두막에서 다른 오두막으로 가는 가장 짧은 길을 따라 가면서 거치는 모든 오솔길들은 항상 산을 내려가는 방향이다.
철수는 등산 마니아이다. 철수가 한 오두막에서 다른 오두막으로 갈 때는 항상 산 정상을 거치는 가장 짧은 길을 따라 간다.
이렇게 간 길의 다양성은 길에 포함된 오솔길의 개수로 정의된다. 두 번 이상 지나간 오솔길은 한 번만 센다는 것에 주의하라.
아래 그림은 가능한 하나의 상황을 보여 준다. 산 정상에 1번 오두막이 있고 3번 오두막과 4번 오두막이 오솔길로 이어져 있다.
아래 그림은 2번 오두막에서 7번 오두막으로 가는 가장 짧은 길을 보여준다.
아래 그림은 2번 오두막에서 7번 오두막으로, 정상을 거쳐서 가는 가장 짧은 길을 보여 준다.
등산로의 구성을 입력으로 받아 모든 가능한 $i$, $j$의 쌍에 대해서 ($1 ≤ i < j ≤ N$), 철수가 $i$번 오두막에서 $j$번 오두막으로 가는 길의 다양성의 총 합을 계산하는 프로그램을 작성하라.
입력
첫 번째 줄에 $N$이 주어진다. 다음 $N−1$개의 줄에 오두막 번호 두 개가 공백 하나를 사이에 두고 주어진다.
두 오두막이 오솔길로 이어져 있다는 뜻이다.
- $2 \leq N \leq 300,000$
출력
첫 번째 줄에 문제의 정답을 출력한다.
풀이
문제 설명이 약간 난해한데, `다양성`은 어떤 경로에서 사용된 간선의 개수를 의미합니다.
루트로 정해진 1번 오두막을 반드시 경유한다는 점을 보면 어떤 오두막 $i$에서 다른 오두막 $j$로 향하는 경로는 $i$ → 1 → $j$입니다. 여기서 모든 간선의 길이는 1이기 때문에 다음과 같은 관찰을 이끌어낼 수 있습니다.
- 정상 1에서 어떤 오두막 $i$까지의 최단거리를 $f(i)$라고 한다.
- 오두막 $i$와 오두막 $j$의 최소공통조상은 오두막 $k$다.
- 오두막 $i$에서 출발하여 오두막 $j$로 도착하는 경로의 다양성은 $f(i) + f(j) - f(k)$다.
처음에는 dfs로 접근해서 직접 각 오두막에서 정상까지의 비용을 따로 계산했었는데 재귀 깊이가 30만까지 되다보니 1번 서브태스크에서 시간 초과가 발생하더라고요. 그렇다고 모든 오두막에 대해 다른 모든 오두막까지의 경로를 직접 계산하기엔 $O(N^2)$의 시간복잡도에 대해 $N$의 범위가 너무 큽니다.
위에서 설명한 3가지를 잘 생각해보면 등산로 전체의 다양성을 빠르게 구할 수 있는 방법이 있습니다. 1~$N$까지의 모든 $i$에 대해 오두막 $i$에서 $i \neq j$인 오두막 $j$까지의 다양성을 정상을 기준으로 `상행`과 `하행`으로 분할해서 구하는 겁니다.
중복되는 경로는 일단 생각하지 않고 보면, 오두막 $i$에서 오두막 $j$까지의 경로에는 반드시 $i$에서 1까지의 경로가 상행이든, 하행이든 포함되게 됩니다. 그러니까 1에서 오두막 $i$까지의 거리를 $N-1$배 하게 되면, 모든 다른 오두막 $j$까지, 또는 $j$로부터의 경로 일부가 계산되는 셈입니다. 이걸 1~$N$에 대해 모두 계산해줍니다.
별로 큰 의미는 없지만 이해를 위해 더 설명해보면 $i$보다 작은 $j$에 대해서는 $i$로 내려오는 하행 경로의 거리가 계산된 것이고 $i$보다 큰 $j$에 대해서는 $i$에서 정상으로 올라가는 상행 경로의 거리가 계산된 셈입니다.
이제 모든 $i$와 $j$에 대해 LCA(최소공통조상) $k$로부터 정상까지의 중복 경로 거리만 제외시키면 됩니다. 이 값은 모든 $k$에 대해 서브트리의 크기를 보면 알 수 있습니다. 이걸 각각의 경로에 대해 LCA를 따로 구해서 1로부터 LCA까지의 거리를 뺀다면 역시 또 시간이 모자라게 됩니다.
현재의 등산로 트리에서 어떤 오두막 $x$를 분리해 $x$가 루트가 되는 서브트리 $T_x$를 만들어봅시다. $T_x$에서 임의의 두 오두막 $a$와 $b$를 선택했을 때, $a$에서 $b$까지의 경로에는 반드시 1부터 $x$까지의 경로가 중복됩니다. $x$가 $LCA(a, b)$가 아니더라도요.
이때, 서브트리의 루트가 그 부모로 연결되는 간선이 얼마나 중복되었는지는 서브트리 내의 정점 2개를 선택하는 경우의 수와 같습니다. 오두막 $x$에 대해 $x$가 루트가 되는 서브트리를 구성하고, 그 서브트리 내의 오두막의 수 $C(x)$를 트리에서의 DP로 쉽게 구할 수 있습니다. 1부터 $N$까지의 모든 $x$에 대해 $Comb(C(x), 2)$를 구해 아까 계산한 모든 상행 경로와 하행 경로 길이의 총합에서 빼주면 됩니다.
정답 코드
import sys
from collections import deque
from heapq import heappush, heappop
input = sys.stdin.readline
def solution():
n = int(input())
link = [[] for _ in range(n+1)]
for _ in range(n-1):
u, v = map(int, input().split())
link[u].append(v)
link[v].append(u)
rank = [-1] * (n+1)
income = [0] * (n+1)
sub = [1] * (n+1)
parent = [i for i in range(n+1)]
rank[1] = 0
bfs = deque([1])
leaf = []
while bfs:
now = bfs.popleft()
for next in link[now]:
if rank[next] == -1:
rank[next] = rank[now]+1
parent[next] = now
bfs.append(next)
income[now] += 1
if income[now] == 0:
heappush(leaf, (-rank[now], now))
while leaf:
_, now = heappop(leaf)
if now == 1: break
sub[parent[now]] += sub[now]
income[parent[now]] -= 1
if income[parent[now]] == 0:
heappush(leaf, (-rank[parent[now]], parent[now]))
res = sum(rank[1:])*(n-1)
for i in range(2, n+1):
res -= sub[i]*(sub[i]-1)//2
print(res)
solution()
