【ABC138 D】「Ki」を解く【Python3】

「AtCoder」解説一覧へ

Nまでの番号が付けられたN個の頂点を持つ根付き木にある条件の操作をしたとき、すべての操作後の各頂点のカウンターの値を求める問題です。

提出
import collections

n, q = map(int, input().split())

ab = [list(map(int,input().split())) for _ in range(n-1)]
px = [list(map(int,input().split())) for _ in range(q)]
cnt = [0]*n

for i in px:
    p, x = i
    cnt[p-1] += x

hen = [[] for _ in range(n)]

for i in ab:
    a, b= i
    hen[b-1].append(a-1)
    hen[a-1].append(b-1)

qu = collections.deque()
qu.append(0)
ans = [-1]*n
ans[0] = cnt[0]

while qu:
    now = qu.popleft()
    for i in hen[now]:
        if ans[i] != -1:
            continue
        ans[i] = ans[now]+cnt[i]
        qu.append(i)

print(*ans)

累積和の進化版のような考え方で解くことができます。

例として、下図のような木を考えていきます。

入力をすべて受け取った後、「cnt」のリストを作ります。

「cnt」には、頂点piに足されたxiを加算していきます。

上図の例では、

cnt = [x1, x2, 0, x3, 0, x4]

となります。

辺でつながっていて、頂点1に近いほうの頂点を親と考えると、子のカウンターの値は「(親のカウンターの値)+(自分に加算されたcnt)」となります。

これを一番親の頂点1から確かめていく必要があります。

そこで「hen」リストを作成し、各頂点がどことつながっているかをまとめます。

上図では、

hen = [[2, 5], [1, 3, 4], [2], [2], [1, 6], [5]]

# リストに加えるときに、a-1, b-1としているため、実際には、数字に「-1」したものが入っている

両端に要素を追加・削除する際に便利な、collectionsモジュールの「deque」を使います。

qu = collections.deque()
qu.append(0)
ans = [-1]*n
ans[0] = cnt[0]

while qu:
    now = qu.popleft()
    for i in hen[now]:
        if ans[i] != -1:
            continue
        ans[i] = ans[now]+cnt[i]
        qu.append(i)

まず、quに0を加え、hen[0]に入っている数字を確認します。

hen[0]には、頂点1とつながっている[2, 5](実際には[1, 4])がはいっています。

ans[1] = ans[0]+cnt[1] # x1+x2
ans[4] = ans[0]+cnt[4] # x1

より、「ans[1]」「ans[4]」が求まります。

quには、[1, 4]が入るので、次のループでhen[1]をチェックします。

hen[1]には[1, 3, 4](実際には[0, 2, 3])が入っており、これも順番に確かめるわけですが、「ans[0]」の値は「-1」ではないので、これを飛ばします。

ans[2] = ans[1]+cnt[2] # x1+x2
ans[3] = ans[1]+cnt[3] # x1+x2+x3

hen[4]には、[1, 6](実際には[0, 5])が入っているため、これをチェックして、

ans[5] = ans[4]+cnt[5] # x1+x4

これまでの処理の後、quには[2, 3, 5]が入っていますが、いずれの場合も

if ans[i] != -1:
    continue

の処理で、ループをスキップしてしまうので、要素の変動はありません。

最終的に

ans = [x1, x1+x2, x1+x2, x1+x2+x3, x1, x1+x4]

となるため、これを空白区切りで出力します。

【ABC138】解説記事リスト

「AtCoder」解説一覧に戻る

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です