Go言語で挑む競プロ #16

前回に引き続きABC138のコンテストに参加したことを書いていこうと思う。
前回のnoteではC問題を解いているので、そちらもどうぞ。

今回は、久々にD問題を解いていこう。
さぁ頑張っていこう。

ABC138Dの問題

問題文
1 から N までの番号がついた N 個の頂点を持つ根付き木が与えられます。
この木の根は頂点 1 で、i 番目の辺 (1≤i≤N−1) は頂点 ai と頂点 bi を結びます。
各頂点にはカウンターが設置されており、はじめすべての頂点のカウンターの値は 0 です。
これから、以下のような Q 回の操作が行われます。
操作 j (1≤j≤Q): 頂点 pj を根とする部分木に含まれるすべての頂点のカウンターの値に xj を足す。
すべての操作のあとの各頂点のカウンターの値を求めてください。

制約
- 2≤N≤2×10^5
- 1≤Q≤2×10^5
- 1≤ai<bi≤N
- 1≤pj≤N
- 1≤xj≤10^4
- 与えられるグラフは木である。
- 入力中の値はすべて整数である。

アルゴリズム

最初は、どう解いたらいいのだろうかと結構頭を悩ませてしまった。いろいろ考えてはみるものも計算時間がネックになりそうでうまい解法が思いつかない。

問題文の通りに愚直にカウンターを足していくと、計算量がO(N*Q)となってしまう。これでは、到底計算時間の制約を守れない...

それで、あーだこーだ考えていたら何とか思いつくことができた。
で、何とか思いついた方法が以下のアルゴリズム。

1. 木を構築する
2. 頂点pjのカウンターにxjを足すことをQ回繰り返す
3. 木の根から葉に向かって以降の探索を始める
4. 現在のノードのカウンターに親ノードのカウンターの値を足す
5. 子ノードへ移動する
6. 全ノードの探索が終了していないなら、手順4へ戻る
7. 各ノードのカウンターの値が解答となる

このアルゴリズムの重要なところは、カウンターを足す処理を後回しにしている点。
pjのカウンターに足すタイミングでは、pjにしか足さず、あとで深さ優先探索を用いて、全ノードのカウンターを更新している。深さ優先探索なら計算量はノード数と枝の数に比例するだけで済む。

もし、最初に考えた愚直にカウンターを足していくアルゴリズムを用いると、pjにカウンターを足すタイミングで、それ以下のノードすべてに足す処理を行う。そのため、すべての操作でpjが木の根を指定された場合、最大でＱ回もＮ個のノードに足す処理が走ることになる。

この差はすごく大きい。実際に必要な計算量の比較を載せてみた。

条件：
　Ｑ＝100, N = 100
毎回カウンター更新：
　O(Q*N)　⇒　10000ループ
最後にカウンター更新：
　O(Q+N)　⇒　200ループ

例として、Q = N = 100としてみても、必要なループ処理の差は9800回！実際の問題では、Q = N = 200000がありうるのだからこの差がいかに大きいかわかると思う。

では、先ほどのアルゴリズムの最終的な計算量はいくらになるのだろう？まとめてみると以下のようになる。

手順1　  ：O(N)
手順2　  ：O(Q)
手順3～6：O(N)
合計　　：O(2N+Q)

全部まとめてもたったこれだけの計算で処理が可能。これなら、計算時間を気にしなくても全然問題ないね。

コード化するために、再度アルゴリズムをまとめてみると。

1. 入力された枝の情報から、木を構成する
2. 入力されるカウンターの値をpjのノードに足す
3. 木の根から深さ優先探索を行い、全ノードを探索する
4. 探索時には、訪れたノードのカウンターの値に親ノードのカウンターの値を足す
5. すべて探索が終了したら、各ノードのカウンターの値を解答とする

全体は難しく感じたけど、一つ一つの手順はとても簡単だね。
さぁこれを実装して、今回も完成だー！

解答

最終的には、以下のようなコードに。

package main

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

type Node struct {
   childs []*Node
   v      int
}

func main() {
   var n, q int
   s := wordScanner(100)
   scanInts(s, &n, &q)

   var a, b int
   l := make(map[int]*Node, n)
   for i := 0; i < n-1; i++ {
       scanInts(s, &a, &b)
       if _, ok := l[b]; !ok {
           l[b] = &Node{[]*Node{}, 0}
       }
       if nodea, ok := l[a]; ok {
           nodea.childs = append(nodea.childs, l[b])
           l[a] = nodea
       } else {
           l[a] = &Node{[]*Node{l[b]}, 0}
       }
   }

   var p, x int
   for i := 0; i < q; i++ {
       scanInts(s, &p, &x)
       l[p].v += x
   }

   dfs(l[1], 0)

   ans := make([]string, n)
   for i := 0; i < n; i++ {
       ans[i] = strconv.Itoa(l[i+1].v)
   }

   fmt.Println(strings.Join(ans, " "))
}

func dfs(node *Node, v int) {
   node.v += v
   for _, child := range node.childs {
       dfs(child, node.v)
   }
}

func wordScanner(n int) *bufio.Scanner {
   s := bufio.NewScanner(os.Stdin)
   s.Split(bufio.ScanWords)
   b := make([]byte, n)
   s.Buffer(b, n)
   return s
}

func scanInts(s *bufio.Scanner, vals ...*int) {
   for i := range vals {
       s.Scan()
       n, _ := strconv.Atoi(s.Text())
       *vals[i] = n
   }
}

いやー今回は、久々にＤ問題を解いたけど、やっぱりＣ問題より難しいね。問題は単純だけども、計算量がネックになってきやすいのかな？

それでも、何とかひらめいて解くことができてよかった。解き終わったときには、すごい達成感だった。
いずれはこれもパパっと解けるようになりたいね。

さぁ次回も頑張っていくぞー！