AtCoder ABC 138 D - Ki (Go) - あるエンジニアのAtCoder奮闘記

irisruneです。休暇を過ごした人もそうでない人もまた心機一転していきましょう。

問題

木構造に慣れていれば単純化しやすい問題ですが、計算量には気を付ける必要があります。

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

var sc *bufio.Scanner

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

type node struct {
    weight int
    adj    []int
    parent int
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, q := nextInt(), nextInt()

    tree := make([]node, n+1)
    for i := 0; i < n-1; i++ {
        a, b := nextInt(), nextInt()
        tree[a].adj = append(tree[a].adj, b)
        tree[b].adj = append(tree[b].adj, a)
    }

    for j := 0; j < q; j++ {
        p, x := nextInt(), nextInt()
        tree[p].weight += x
    }

    ans := make([]int, n+1)
    queue := make([]int, 0)
    queue = append(queue, 1)
    for len(queue) > 0 {
        u := queue[0]
        queue = queue[1:]
        ans[u] = ans[tree[u].parent] + tree[u].weight
        for _, v := range tree[u].adj {
            if v == tree[u].parent {
                continue
            }
            tree[v].parent = u
            queue = append(queue, v)
        }
    }

    for i := 1; i < n; i++ {
        fmt.Print(ans[i])
        fmt.Print(" ")
    }
    fmt.Println(ans[n])
}

木をどのように構築するかは様々な方法があると思いますが、今回は頂点を構造体で定義して隣接リストをもつ形を取りました。

まず愚直に解くことを考えてみると、各操作を行うたびに各頂点のカウンターに値を足すことになります。しかし、操作回数が $Q$ 回、各操作の対象となる頂点数が $O(N)$ であることから、計算量が $O(NQ)$ となってしまいます。

簡単に思い付く計算量削減の方法としては、各操作を基準となる頂点ごとにまとめてしまうことです。これにより操作にかかる計算量が $O(Q)$ となり全体の計算量が $O(N+Q)$ で抑えられるように見えます。

ここで重要なのが各操作をまとめたものを木に対してどのように反映させるかです。これまた愚直に実装すると各頂点 $1,2,\ldots,N$ にまとめた操作の分量だけ、頂点を根とする部分木の各頂点に対してカウンターの値を増やすという形になりますが、これでは計算量が $O(N^2+Q)$ となってしまいます。