CodeChef October Challenge 2019: Bacterial Reproduction

CodeChef October Challenge 2019の問題: Bacterial Reproduction (Code: BACREP)

問題ページ: Contest Page | CodeChef

問題概要

根頂点が頂点1の根付き木がある。

はじめ、各頂点に細菌が存在し、頂点には匹存在する。

各秒において、以下が発生する。

• 葉でない各頂点について、子頂点の数をとする。その頂点に存在する細菌を匹とした時、匹に分離した上で各子頂点に匹ずつ移動する。(この時、その頂点に匹の細菌はいなくなる)
• ある1頂点に細菌が何匹か出現することがある

秒目から秒目 の 秒間の各秒に対し、以下のクエリを処理せよ

• '+ v k': この秒の間に頂点に匹の細菌が出現する (ただしこの秒数では分裂移動しない)
• '? v': この秒で細菌が分離し移動した時点での頂点の細菌の数を出力する

制約

• 
• 
• 

解法

オフラインでクエリを処理する。

長さの配列を用意し、DFSで各頂点を探索しながら配列を更新し、その頂点の各時刻の細菌の数を求められるようにする。
このDFSで頂点を移動する際には配列の時刻0の基準をずらすようにする。具体的には、深さの頂点における時刻の細菌の数がになるように管理する。
ただし、葉となる頂点に到達した細菌は移動しないため、時刻の細菌の数は時刻までの累積になる。

この配列を更新しながら、根付き木を根からDFSで探索しながら各解を計算していく。

DFSの具体的な処理として、根頂点1から順に、深さの頂点について以下のように処理していく。

1. に を加算
2. 頂点に秒目に匹出現する場合に配列の に を加算
3. 頂点の秒目の細菌の数を出力する場合、以下の値を求める
   • 頂点が葉の場合、 ~ の総和
   • 頂点が葉でない場合、配列の値
4. 子頂点を探索
5. 親頂点に戻る場合は(1), (2)で加算した値を戻す

問題のサンプル入力における、各頂点を訪れた際の配列は以下のようになる。

サンプル入力の処理 (丸内の数字は頂点番号)
各配列について、灰色はその頂点の初期細菌、赤色はその頂点に出現した細菌、緑色はクエリで答えるべき時刻、下の小さい数字は時刻を表す

総和を求めるため、配列はBITとして管理すればよい。

計算量は

実装

提出コード(C++14): Solution: 27000541 | CodeChef

#define N 500005

int n;
vector<int> g[N];

#define D 1000005

int d;
ll data[D+1];
void add(int k, ll x) {
  while(k <= d) {
    data[k] += x;
    k += k & -k;
  }
}
ll get(int k) {
  ll s = 0;
  while(k) {
    s += data[k];
    k -= k & -k;
  }
  return s;
}

vector<P> ps[N];
vector<int> qs[N];
ll ans[N];
ll a[N];

int dfs(int v, int p, int d) {
  // 細菌の追加
  add(n-d, a[v]);
  repe(&p, ps[v]) {
    int j = p.first, k = p.second;
    add(n-d+1+j, k);
  }
  // 解の計算
  if((p != -1 && g[v].size() == 1) || (p == -1 && g[v].size() == 0)) {
    repe(j, qs[v]) {
      ans[j] = get(n-d+1+j);
    }
  } else {
    repe(j, qs[v]) {
      ans[j] = get(n-d+1+j) - get(n-d+j);
    }
  }
  // 子の探索
  repe(w, g[v]) {
    if(w == p) continue;
    dfs(w, v, d+1);
  }
  // 細菌の数を戻す
  add(n-d, -a[v]);
  repe(&p, ps[v]) {
    int j = p.first, k = p.second;
    add(n-d+1+j, -k);
  }
  return 0;
}

int main() {
  int q;
  cin >> n >> q;
  d = n+q;
  rep(i, n-1) {
    int x, y; cin >> x >> y; --x; --y;
    g[x].push_back(y);
    g[y].push_back(x);
  }
  rep(i, n) cin >> a[i];
  rep(i, q) {
    char c; cin >> c;
    int v, k;
    switch(c) {
      case '+':
        cin >> v >> k; --v;
        ps[v].emplace_back(P{i, k});
        break;
      case '?':
        cin >> v; --v;
        qs[v].push_back(i);
        break;
      default:
        assert(0);
        break;
    }
    ans[i] = -1;
  }
  dfs(0, -1, 0);
  rep(i, q) {
    if(ans[i] != -1) {
      cout << ans[i] << "\n";
    }
  }
  return 0;
}