記事を書くに至った背景
競技プログラミングを始めて、数ヶ月が経ちますが、
これまでは基礎的なアルゴリズムをしっかり理解することを怠っていました。
今回つまづいたDFSに関しても、概要だけ調べただけだったので、
======
後輩「先輩、DFSってなんですか？」
僕 「DFS? 深さ優先探索のことだよ」
後輩「深さ優先探索ってなんですか？」
僕 「えっ、深さ優先探索？DFSのことだよ」
後輩「...」
僕 「...」
=======
状態でした。
そして、AtCoder Beginner Contest 138の終了後、
今回解けなかったD問題の解説を読みました。
あんまりよく分からなかったので、
他の人のACしているコードを見ることにしました。
何人かのACコードを見ていると、dfs()と名付けられた関数を定義しているコードが多いことに気づき、
僕 「あー、DFSを使えば解けたのか（？）」
ということで、
DFSのことを、ろくに分かっていないくせに
もう理解した気になっていました。
しかし、
後から始めた同期や後輩に
レートが抜かれそうになっている現状への焦りもあり、
このままではいけないと感じ、
この機会に、DFSというものを理解しようと心に決めました。
これまでサボってきたつけもあり、結果的にDFSを理解して、
ACを通すのに約4時間かかったのですが、
なんだかんだで理解が深まったので、
この問題でつまづいてる他の人が
これを理解するのに僕のように４時間もかけずに済むように、
記事を書いてみようと思った次第です。
対象の問題
前書きが長くなりましたが、本題に移ります。
https://atcoder.jp/contests/abc138/tasks/abc138_d
木が与えられた後、ノードと数字(p,x)が複数与えられるので、
ノード以下の部分木に数字xを足していくといったものでした。
考え方
迷走しましたが、やること自体は、複雑ではありませんでした。
例えば、下のような木を考えます。
この木に対して、次のようなノードと数字の組み合わせが入力されるとします。
(p, x)
(1, 100)
(5, 10)
(8, 5)
入力データとしては、次のようになります。
8 3
1 2
1 3
2 4
3 5
3 6
5 7
5 8
1 100
5 10
8 5
まず、(1,100)ですが、
1のノードに100を割り当てます。（実装には配列を使う)
次は、(5,10)なので、
5のノードに10を割り当てます。
そして、最後は(8,5)なので、
8のノードに5を割り当てます。
割り当てがない部分は0と考えます。
そしたら、
「子ノードに、自分の持つ数字を足す」という処理を
一番てっぺんに当たる根のノードから、再帰的に行なっていきます。
ここがDFSに当たる部分でした。
図で表すとこんな感じです。↓
最終的には、こうなります。
このキモの部分をいわゆるDFS（深さ優先探索）で行うわけですが、
このアルゴリズム自体の説明は、
僕が説明するよりもっといい記事がいくらでもあるので、そちらを見ると良いと思います。
僕が分かりやすいと思ったのは、
https://www.slideshare.net/chokudai/dfs-49066641
です。
過去のAtCoderコンテストの解説スライドです。
さすがのAtCoderですね。
コード
最後にpythonで実装したコードも載せておきます。
import sys
input = sys.stdin.readline #for input speed
sys.setrecursionlimit(10**6) #for deep recursion
n, q = map(int, input().split())
tree_index = [[] for _ in range(n)]
for i in range(n-1):
a,b = map(int, input().split())
a -= 1
b -= 1
tree_index[a].append(b) # a<bとは限らないため、双方向にする
tree_index[b].append(a)
ans = [0] * n
for i in range(q):
p, x = map(int, input().split())
ans[p-1] += x
def dfs(cur, par):
for chl in tree_index[cur]:
if chl == par: #parentとchildが等しい時はスキップ
continue
ans[chl] += ans[cur]
dfs(chl, cur)
dfs(0, -1)
print(*ans)
木構造を表すために、tree_indexで配列を作っていますが、
↓の記事にもあるように、a<bに限らないという制約が追加？されたので、
辺に対してaとbのノードが与えられたときに、
双方向に対して、関係性を持たせています。
AtCoder Beginner Contest 138 D - Ki のテストケースがコンテスト後に増えている件について (修正版)
感想
やっぱ適当じゃなくて、腑に落ちるまで復習した方がいいと実感しました。