Segtreeのテクニック

昨日のCF259Div1E問題で新しいテクニックを知ったのでメモします。

とりあえずまず遅延伝播から書こう。

遅延伝播

以下のクエリを処理せよ
Addクエリ:区間内の全要素に1を足す
Sumクエリ:区間の和を求める

という問題を考える。
ノードに [ 区間の和 ] を持つsegtreeを作ると、SumクエリはlogNで処理できる。
しかし、Addクエリでは更新しなければならないノードが大量発生してしまう。
Addクエリを高速化したいので、Addクエリのときの操作を「ノードの [ 区間全体に1を足された回数 ] をインクリメントする」とする。
そして、Sumクエリの際に訪れたノードのみに関してそれを伝播させて行く。つまり、
・ノード i の [ 区間の和 ] → tot[i]
・ノード i の [ 区間全体に1を足された回数 ] → cnt[i]
とすると、
・tot[i] += cnt[i] * [ノード i に対応する区間の長さ]
・cnt[i<<1] += cnt[i]
・cnt[i<<1|1] += cnt[i]
・cnt[i] = 0
という操作をすれば良いことになる。(Addクエリの時はtotも少し更新する必要がある)
すると、どちらのクエリもlog Nで処理できるようになる。

#include <cstdio>
#include <vector>
#include <algorithm>
#include <iostream>
using namespace std;

struct seg {
  vector<int> tot, cnt;
  int sz;
  seg() {}
  seg(int n) {
    sz = 1;
    while (sz < n) sz <<= 1;
    tot.resize(sz << 1);
    cnt.resize(sz << 1);
  }
  void add(int a, int b, int i=1, int l=0, int r=-1) {
    if (r == -1) r = sz;
    if (a <= l && r <= b) { ++cnt[i]; return;}
    tot[i] += min(r, b) - max(l, a);
    int c = l+r>>1;
    if (a < c) add(a,b,i<<1,l,c);
    if (c < b) add(a,b,i<<1|1,c,r);
  }
  int sum(int a, int b, int i=1, int l=0, int r=-1) {
    if (r == -1) r = sz;
    tot[i] += (r-l) * cnt[i];
    if (i < sz) {
      cnt[i<<1] += cnt[i];
      cnt[i<<1|1] += cnt[i];
    }
    cnt[i] = 0;
    if (a <= l && r <= b) return tot[i];
    int c = l+r>>1, res = 0;
    if (a < c) res += sum(a,b,i<<1,l,c);
    if (c < b) res += sum(a,b,i<<1|1,c,r);
    return res;
  }
};

int N = 8;
vector<int> type, l, r;
void solve() {
  seg d(N);
  for (int i = 0; i < (int)type.size(); ++i) {
    if (type[i] == 0) { // Add
      d.add(l[i], r[i]);
    } else { // Sum
      cout << d.sum(l[i], r[i]) << endl;
    }
  }
}

void addQuery(int ntype, int nl, int nr) {
  type.push_back(ntype);
  l.push_back(nl);
  r.push_back(nr);
}

int main(){
  addQuery(0, 0, 8); // 1 1 1 1 1 1 1 1
  addQuery(1, 0, 8); // 8
  addQuery(0, 1, 5); // 1 2 2 2 2 1 1 1
  addQuery(0, 4, 6); // 1 2 2 2 3 2 1 1
  addQuery(1, 3, 7); // 8
  addQuery(0, 2, 3); // 1 2 3 2 3 2 1 1
  addQuery(0, 7, 8); // 1 2 3 2 3 2 1 2
  addQuery(1, 0, 8); // 16
  addQuery(1, 1, 7); // 13
  solve();
  
  return 0;
}

というのが遅延伝播のテクニックです。
区間全体への操作を一番上の部分にだけとりあえず記録しておき、実際に値を参照する時にそれを伝播すると言う感じです。

本題

E問題は以下のような問題が解ければ解ける。

N個のコップが並んでいて、毎秒1ccずつ水が溜まる。
コップ i の容量は Mi であり、それ以上の水が来ると溢れてどこかへ消える。
以下のクエリをQ個処理せよ。
クエリ:時間tに区間の水を全部捨てる。捨てた水の合計量を答えよ。

違う書き方をするとこうなる。(そのあたり詳しくはhogloidのブログを見て下さい)

0初期化された配列Aと、配列の各要素について関数Fが与えられる。
以下のクエリを処理せよ。
・区間と整数tが与えられるので、ΣF_i(t-A[i])を求め、区間内のA[i]を全てtに書き換える(tはクエリごとに単調増加)
ただし、任意の区間についてO(区間長 log 区間長)の前計算をすることによりΣF_i(x)をO(log 区間長)で求めることが出来るようになる。

要するに、区間内のAの値が全て同じなら、その区間についての値をO(log len)で求めることが出来る。
とりあえずsegtreeの各ノードについて、対応する区間について前処理をしておいて、各ノードに関して「区間内のAの値が全て同じなら、その区間についての値をO(log len)で求めることが出来る」という状況を作っておく。
そして、各ノードに関してTという値(名前は何でも良いけど)を持っておく。
クエリが飛んで来たときは、
・T[i<<1] = max(T[i<<1],T[i]), T[i<<1|1] = max(T[i<<1|1],T[i]) と伝播させる
・ノード i がクエリの区間に完全に含まれていてかつT[i] != -1ならΣF_i(t-T[i])を返す
・そうでないなら範囲内のところに潜り、帰って来たら、
 ・ノード i がクエリの区間に完全に含まれているならT[i] = tとし、
 ・そうでないならT[i] = -1とする
というふうに処理をする。
すると、クエリ1回につき-1はO(logN)個しか増えないので、ならし計算量O(N log^2 N)みたいな感じになる。
よく分からなければ、実際にE問題を解いてみるか、ソースコードを読んでみると良いでしょう。(ソースコードでは-1ではなく0にしてます。)
初期値が定まっていてもちょっとソースコードを変えるだけで対応できるあたりがクール。


奈良市計算量がうまく決まっているので、個人的にこのテクニックを奈良木と呼ぶことにした。

追記:今はsegtree beatsって呼ばれてそう