F - Distance Sums 2

[hxrxchang]

https://atcoder.jp/contests/abc220/tasks/abc220_f

問題概要

頂点がN個の木がある。 1~Nの各頂点について、その頂点からすべての頂点へ行く場合に辿る辺の数の合計を求めよ。

解き方

• 一回目のdfsで "頂点0をルートとした木の中で、各頂点を根とする部分木の距離の総和" を求める。
• 二回目のdfsで "各頂点から他のすべての頂点へ行く距離の総和" を求める。

ことができているらしい...

var n int
var tree [][]int
var subTreeSize []int
var distanceSum []int

func solve() {
    n = getInt()
    tree = make([][]int, n)
    for i := 0; i < n-1; i++ {
        in := getInts()
        a, b := in[0]-1, in[1]-1
        tree[a] = append(tree[a], b)
        tree[b] = append(tree[b], a)
    }

    // subTreeSize[i]: iを根とする部分木のサイズ
    subTreeSize = make([]int, n)
    distanceSum = make([]int, n)

    dfs1(0, -1)

    dfs2(0, -1)

    for i := 0; i < n; i++ {
        fmt.Println(distanceSum[i])
    }
}

// dfs1を呼び出し後、 distaneSum[i] は 0をルートとした木の中でiを根とする部分木の距離の総和になる
func dfs1(node, parent int) {
    subTreeSize[node] = 1
    distanceSum[node] = 0
    for _, next := range tree[node] {
        if next == parent {
            continue
        }
        dfs1(next, node)

        subTreeSize[node] += subTreeSize[next]
        distanceSum[node] += distanceSum[next] + subTreeSize[next]
    }
}

// dfs2を呼び出し後、 distaneSum[i] は 頂点iから他の頂点への距離の総和になる
func dfs2(node, parent int) {
    for _, next := range tree[node] {
        if next == parent {
            continue
        }
        distanceSum[next] = distanceSum[node] - subTreeSize[next] + n - subTreeSize[next]
        dfs2(next, node)
    }
}

https://atcoder.jp/contests/abc220/submissions/53730924