木dpについて
- 木構造のdp
- 典型的には
dfsの中でdpを管理する
すべての頂点の最短距離和を計算するには貢献度を計算する
問題
解説
- すべての点を求めるにはワーシャルフロイド法やベルマンフォード法が必要になるが、
O(n^2)で計算できない - ある辺から通る数を管理することで貢献度という指標を計算する
- この数(r)は、すべての最短距離の和の一部
r*(N-r)になる - すべての辺のrに対して和を取ることで、すべての
最短距離の和になる
- この数(r)は、すべての最短距離の和の一部
解答
import sys; sys.setrecursionlimit(10**9)
N=int(input())
G=[[] for _ in range(N)]
AB = []
for _ in range(N-1):
a,b=map(int,input().split())
a-=1;b-=1
AB.append( (a,b))
G[a].append(b)
G[b].append(a)
dp = [0]*N
def dfs(pos, prv):
dp[pos] = 1
for i in G[pos]:
if i == prv:
continue
dfs(i, pos)
dp[pos] += dp[i]
dfs(0,-1)
# print(dp) # 貢献度のリスト
ans = 0
for i in range(N-1):
a,b=AB[i]
r=min(dp[a], dp[b])
ans += r*(N-r)
print(ans)