競プロ初中級者100問

ABC138D Ki

この記事で使うアルゴリズム

深さ優先探索

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

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

この問題は深さ優先探索を用いて解くことができます。

深さ優先探索はスタックや再帰関数によって実装できます。今回は再帰関数を用います。

方針

  • 操作の際の根に当たるノードのカウンターVに先に値を加算する。
  • 深さ優先探索を行いながら、親ノードのカウンターの値を加算していく。
    • 深さ優先探索ではlinkを参照してノードを追加します。
      • この時、親ノードを追加して無限ループとならないように除外します。そのためにparentで親ノードを記憶しています。
  • 最後にカウンターVの値を文字列結合して出力する。

ポイント

入力行数が多いので高速化が必要です。

1
input = sys.stdin.readline

デフォルト値だと小さすぎるので、再帰処理の上限を引き上げる。

1
sys.setrecursionlimit(int(1E+7))

コード（全体）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import sys
input = sys.stdin.readline
sys.setrecursionlimit(int(1E+7))
def main():
    N, Q = map(int, input().split())
    link = [[] for _ in range(N+1)]
    for n in range(N-1):
        a, b = map(int, input().split())
        link[a].append(b)
        link[b].append(a)
    V = [0] * (N+1)
    for q in range(Q):
        p, x = map(int, input().split())
        V[p] += x # `操作`の際の根に当たるノードのカウンター`V`に先に値を加算する。

    def dfs(i, parent, acc): # 深さ優先探索
        V[i] += acc # 親ノードまでの累積値を加算
        for j in link[i]:
            if j != parent: # 親ノードを追加しない
                dfs(j,i,V[i])
    cur = 1
    parent = 0
    acc = 0
    dfs(cur, parent , acc) # 親ノードを0としてスタート
    V = [str(v) for v in V]
    print (" ".join(V[1:]))

if __name__ == '__main__':
    main()

記事情報

  • 投稿日：2020年4月15日
  • 最終更新日：2020年5月8日