けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 138 D - Ki (緑色, 400 点)

まさに「the 緑 diff」な問題だと思う。
「木を上手く扱えるか」を問いかける問題。

問題へのリンク

問題概要

 N 頂点の木が与えられる。この木の頂点 1 を根とした根付き木を考える。各頂点には初期状態では 0 という数値が書かれている。以下の  Q 回の操作を行なった後の各頂点の数値を答えよ。

  •  j 回目の操作では、頂点  p_{j} を根とする根付き木に含まれるすべての頂点に  x_{j} を加算する

制約

  •  N, Q \le 2 \times 10^{5}

考えたこと

この問題は、2 つのことを要求している

  • 木じゃなくて配列なら解けるのか?
  • 木はどう扱えばよいのか?

これらそれぞれ単体では、茶色 diff になるのかなと思う。

木じゃなくて配列ならどうか?

この問題に限らず、木上のクエリを処理する問題では、「まずは配列について考える」というのがポイントになることが多い。配列は「枝分かれのない木」なので、木の特殊ケースとみなせるのだ。

さて、今回の問題を木に置き換えると、次のような問題になる。


  •  N 要素からなる配列があって、最初はすべて 0 が書き込まれている
  • 各クエリは「配列の  p, p + 1, \dots, N に、値  x を加算する」

これは各クエリに対して愚直に計算すると  O(NQ) になってしまうのでダメ。

ここで「いもす法」のようなアイディアを思い出してみよう。まず、「配列の  p, p + 1, \dots, N に、値  x を加算する」という操作は次の 2 段階の操作に分解できる。

  1. 配列の index p の位置に、 x を加算する
  2. 配列の累積和をとる

通常は、各クエリごとに、「1」をやって「2」をやるという感じのイメージだけど、

  • まず 1 について、全クエリ分を処理してしまう (それぞれのクエリにつき  O(1))
  • 最後に一回だけ累積和をとる ( O(N))。

という風にすれば、 O(N + Q) で処理できる。これはいもす法の特殊ケースで、「一律加算する右側が常に配列の右端になっているような場合」とみなせる。

木上で

これまでの処理を自然に木に拡張すれば OK。配列上で累積和をとる操作は、木では「根から子へと値を配っていく」というイメージになる。

まず配列 res を 0 に初期化しておく。そして、各クエリ ( p, x) ごとに、res の index  p のところに  x を加算しておく。その後、木の根から出発して、子へと値を累積させていく。最終的には各頂点 v に対して、

  • res[ v ] = 根から頂点 v に至るまでの各頂点までの  x の値の総和

となるようにすれば OK。以下のような再帰関数を定義してあげて、

// 頂点 v について処理、p: v の親
dfs(int v, int p, vector<long long> &res);

次のように処理することで実現できる。

  • val[ v ] += val[ p ] (v が根の場合はしない)

計算量は  O(N + Q)

#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;
}