-
Notifications
You must be signed in to change notification settings - Fork 530
/
d.go
40 lines (37 loc) · 988 Bytes
/
d.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
package main
// github.com/EndlessCheng/codeforces-go
func sumOfDistancesInTree(n int, edges [][]int) []int {
g := make([][]int, n)
for _, e := range edges {
v, w := e[0], e[1]
g[v] = append(g[v], w)
g[w] = append(g[w], v)
}
sz := make([]int, n)
var f func(int, int) int
f = func(v, fa int) (sum int) { // sum 表示以 0 为根的子树 v 中的节点到 v 的距离之和
sz[v] = 1
for _, w := range g[v] {
if w != fa {
sum += f(w, v) + sz[w] // 子树 w 的每个节点都要经过 w-v,因此这条边对 sum 产生的贡献为 sz[w]
sz[v] += sz[w]
}
}
return
}
sum0 := f(0, -1)
ans := make([]int, n)
var reroot func(v, fa, sum int)
reroot = func(v, fa, sum int) {
ans[v] = sum
for _, w := range g[v] {
if w != fa {
// 离子树 w 中的所有节点近了 1,又离不在子树 w 中的节点远了 1
// 所以要减去 sz[w],并加上 n-sz[w]
reroot(w, v, sum+n-sz[w]*2)
}
}
}
reroot(0, -1, sum0)
return ans
}