まさに「the 緑 diff」な問題だと思う。
「木を上手く扱えるか」を問いかける問題。
問題概要
頂点の木が与えられる。この木の頂点 1 を根とした根付き木を考える。各頂点には初期状態では 0 という数値が書かれている。以下の 回の操作を行なった後の各頂点の数値を答えよ。
- 回目の操作では、頂点 を根とする根付き木に含まれるすべての頂点に を加算する
制約
考えたこと
この問題は、2 つのことを要求している
- 木じゃなくて配列なら解けるのか?
- 木はどう扱えばよいのか?
これらそれぞれ単体では、茶色 diff になるのかなと思う。
木じゃなくて配列ならどうか?
この問題に限らず、木上のクエリを処理する問題では、「まずは配列について考える」というのがポイントになることが多い。配列は「枝分かれのない木」なので、木の特殊ケースとみなせるのだ。
さて、今回の問題を木に置き換えると、次のような問題になる。
- 要素からなる配列があって、最初はすべて 0 が書き込まれている
- 各クエリは「配列の に、値 を加算する」
これは各クエリに対して愚直に計算すると になってしまうのでダメ。
ここで「いもす法」のようなアイディアを思い出してみよう。まず、「配列の に、値 を加算する」という操作は次の 2 段階の操作に分解できる。
- 配列の index p の位置に、 を加算する
- 配列の累積和をとる
通常は、各クエリごとに、「1」をやって「2」をやるという感じのイメージだけど、
- まず 1 について、全クエリ分を処理してしまう (それぞれのクエリにつき )
- 最後に一回だけ累積和をとる ()。
という風にすれば、 で処理できる。これはいもす法の特殊ケースで、「一律加算する右側が常に配列の右端になっているような場合」とみなせる。
木上で
これまでの処理を自然に木に拡張すれば OK。配列上で累積和をとる操作は、木では「根から子へと値を配っていく」というイメージになる。
まず配列 res を 0 に初期化しておく。そして、各クエリ () ごとに、res の index のところに を加算しておく。その後、木の根から出発して、子へと値を累積させていく。最終的には各頂点 v に対して、
- res[ v ] = 根から頂点 v に至るまでの各頂点までの の値の総和
となるようにすれば OK。以下のような再帰関数を定義してあげて、
// 頂点 v について処理、p: v の親 dfs(int v, int p, vector<long long> &res);
次のように処理することで実現できる。
- val[ v ] += val[ p ] (v が根の場合はしない)
計算量は 。
#include <bits/stdc++.h> using namespace std; using Graph = vector<vector<int>>; int N, Q; Graph G; void dfs(int v, int p, vector<long long> &res) { if (p != -1) res[v] += res[p]; for (auto e : G[v]) { if (e == p) continue; dfs(e, v, res); } } int main() { cin >> N >> Q; G.assign(N, vector<int>()); for (int i = 0; i < N-1; ++i) { int a, b; cin >> a >> b; --a, --b; G[a].push_back(b); G[b].push_back(a); } vector<long long> val(N, 0); for (int i = 0; i < Q; ++i) { int p, x; cin >> p >> x; --p; val[p] += x; } dfs(0, -1, val); for (auto v : val) cout << v << " "; cout << endl; }