AtCoder Beginner Contest 138 D - Ki

お題箱より。

D - Ki

様々な木の問題を解く上で基本となるような実装なので、実装も含めて解説します。

解法

1回の操作のイメージを図にすると以下のようになります。

f:id:betrue12:20200224171708p:plain

実際には大量の操作があるので、これを1つ1つ処理していくことはできません。ではどうするかと言うと、以下の図のように考えます。

f:id:betrue12:20200224172627p:plain

まず与えられた $Q$ 回の操作について、頂点 $p_{i}$ 「だけ」に値 $x_{i}$ を足していきます。それが終わったら、根から順番に足された値を「押し流していく」ように伝えていきます。

図中の $p_{1}, p_{2}$ のように、伝えていった先に既に別の値がある場合、そのまま足し算をすれば良いです。このようにすると、 $x_{1}, x_{2}, x_{3}$ がちゃんと $p_{1}, p_{2}, p_{3}$ を根とする部分木に伝わっていることが図から分かります。

実装

これを実装しましょう。木において「根から順番に葉まで」あるいは「葉から順番に根まで」値を伝えながら処理をしていきたい場合は、DFSをすることが多いです。

まずは実装の全体像を示します。

#include <bits/stdc++.h>
using namespace std;

int N, Q;
vector<int> edges[200000];
int64_t num[200000];

void dfs(int i, int p){
    for(int j : edges[i]) if(j != p){
        num[j] += num[i];
        dfs(j, i);
    }
}

int main(){
    cin >> N >> Q;
    for(int i=0; i<N-1; i++){
        int a, b;
        cin >> a >> b;
        edges[a-1].push_back(b-1);
        edges[b-1].push_back(a-1);
    }
    for(int i=0; i<Q; i++){
        int p, x;
        cin >> p >> x;
        num[p-1] += x;
    }
    dfs(0, -1);
    for(int i=0; i<N; i++) cout << num[i] << " \n"[i==N-1];
    return 0;
}

まずは入力を受け取ってグラフを構成します。edges[i] には頂点 i と辺で接続されている頂点の番号が入っています。このようにグラフを記憶する方法を「隣接リスト」と呼びます。

また変数 num に、各操作の「頂点 $p_{i}$ だけに値 $x_{i}$ を足す」という処理を行います。

そしてDFSをします。DFS部分のコードを再度抜き出します。

void dfs(int i, int p){
    for(int j : edges[i]) if(j != p){
        num[j] += num[i];
        dfs(j, i);
    }
}

引数 i は今見ている頂点、p はその親です。今回グラフを無向グラフとして扱っていて、edges[i] の中には親 p も含まれているので、下に下に見ていく際には親のほうに戻らないようにしないといけません。そのため p を引数として渡し、if(j != p) の条件で弾いています。

i の隣接頂点 j を見ていって、もし j が親でない（つまり子である）ならば、まずは値を伝えます（num[j] += num[i];）。そして再帰的に、頂点 $j$ について dfs 関数を呼び出します。

こうすることで、先ほどの説明で書いた「根から順番に足された値を押し流していく」という処理が実現できることになります。

あとは、細かい補足をしておきます。

dfs関数を最初に呼び出すときは根を引数にしますが、根に親はないので、他の頂点番号のどれとも異なる値をセットしておけば良いです。私はよく -1 を使っています。
dfs関数内部の for(int j : edges[i]) は範囲for文と呼ばれるもので、vector などの各要素を順番に見ていくようなループを回します。便利。
今回の問題の場合は「根をスタートにして、頂点を飛び越さないような順番で見ていく」限りはどのような順番で見ても良いので、DFSじゃなくてBFSでも正しい答えを求めることができます。ただDFSのほうが再帰で書きやすいことが多く、根付き木の頂点を順番に処理していく問題はほとんどの場合DFSがオススメです。