D - Ki

[hxrxchang]

https://atcoder.jp/contests/abc138/tasks/abc138_d

問題概要

木がある。 それぞれの頂点がカウンターを持っており、初期値は0になっている。 Qのクエリが与えられるので、その度に以下の操作をする。

• p, xが与えられる。頂点pをルートとした部分木の全要素のカウンターに、xを足す。

最終的な全要素のカウンターの状態を求めよ。

解き方

制約は N, Qともにmax 2 * 10 * 5 クエリの度にDFSで更新していくと、O(N Q) で間に合わない。 そこで累積和の考え方を使って O(N) で済む方法を考える。 部分木全体を更新することから、 ある頂点を考えたとき、その頂点のカウンターは親に依存していることが分かる。 そのため最終的な親のカウンターの数値だけ分かれば、自分のカウンターは親の数値に自分に対するクエリがあればそれを足すことで求められる。 頂点1がルート(親がない = 依存がない)なので、頂点1の最終的なカウンターの数値はQから分かる。 なので1からDFSで子を辿って答えを求められる。

var tree [][]int
var dp []int

func solve() {
    in := getInts()
    n, q := in[0], in[1]

    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)
    }

    dp = make([]int, n)

    for i := 0; i < q; i++ {
        in := getInts()
        p, x := in[0]-1, in[1]
        dp[p] += x
    }

    dfs(0, -1)

    ans := make([]string, n)
    for i := 0; i < n; i++ {
        ans[i] = i2s(dp[i])
    }
    fmt.Println(strings.Join(ans, " "))
}

func dfs(node, parent int) {
    for _, next := range tree[node] {
        if next == parent {
            continue
        }
        dp[next] += dp[node]
        dfs(next, node)
    }
}

https://atcoder.jp/contests/abc138/submissions/53681159