セグ木は知っているとします。 知らない人は えびちゃんのスライド なり適当な記事なりを読んでください。
セグ木上のにぶたんも上のスライドに載っていますが、もう少し解説したいかなと思います。 上のスライドにある図を見るとイメージが湧きやすいかもしれません。
愚直にやると、セグ木の区間和*1取得と、二分探索で、ふたつ log がついてしまいます。 セグ木の実装を非再帰でやって、\([l, r)\) の区間和を \(O(\log(r- l))\) でできるとしても、オーダーは落ちません。
参考までに、\(n = 10^6\) とすると、\(n\log(n)^2\) は \(n\sqrt{n}/2.5\) 程度ですが、\(n\log(n)\) は \(n\sqrt{n}/50\) 程度です。 落とせる log は落としておきたい気がします。
この方針以外の実装もある気がしますが、えびちゃん的にはこれが好きなので、とりあえずこれを紹介します。 気づいたらそのうち宗派が変わっているかもしれません。
定義
セグ木の葉の個数を \(n\) とします。 セグ木の元の配列長が \(n\) ということです。 区間 \([l, r)\) の和を \(a[l, r)\) と書くことにします。
左端を固定するとき
以下を満たす述語 \(p\) が与えられるとします。
- 単位元 \(e\) に対して \(p(e)\) が true
- ある \(i\) が存在して \(p(a[l, i))\) が false であれば、任意の \(i \lt j \le n\) に対して \(p(a[l, j))\) が false
このとき、\(p(a[l, i))\) が true かつ、次のいずれかを満たす \(i\) を返します。
- \(i = n\)
- \(i < n\) かつ \(p(a[l, i+1))\) が false
右端を固定するとき
以下を満たす述語 \(p\) が与えられるとします。
- 単位元 \(e\) に対して \(p(e)\) が true
- ある \(i\) が存在して \(p(a[i, r))\) が false であれば、任意の \(0 \le j \lt i\) に対して \(p(a[j, r))\) が false
このとき、\(p(a[i, r))\) が true かつ、次のいずれかを満たす \(i\) を返します。
- \(i = 0\)
- \(i \gt 0\) かつ \(p(a[i-1, r))\) が false
補足
一度 false になると、区間を広げても false であり続ける述語を受け取ります。 条件を満たす最大の区間を求めているのに相当します。 両方の場合において、引数と返り値を昇順に \(l, r\) とすると、\(p(a[l, r))\) は true となります。
\(p(e)\) が true であることを要求しているのは、そうしないと返り値がつらくなるためです。
-1 を返すのはあまりうれしくないし、そもそも単位元の時点で false なら、呼び出す前から答えがわかりますし。
なお、これは ACL のものの定義と同じです。
簡単な場合
まずは、セグ木を完全二分木に限定し、二分探索の始点を左端に固定した場合を考えます。 全区間の和を述語に渡した場合は false になるとします（そうかどうかの判定は述語を高々一回呼べばわかる）。
これは、必要な情報を管理しながら、ノードを下っていくことでできます。 必要な情報というのは、「今まで足されている値」「今見ている添字」くらいだと思います。
適当に非再帰で簡単に実装できます。
たとえば、左に潜るなら i <<= 1 のようなことをすればよいです。
auto x = e; // 今まで足した値。単位元で初期化 while (i < n) { i <<= 1; auto tmp = op(x, a[i]); if (p(tmp)) { // 左の子を足しても大丈夫 x = tmp; // 足した結果で上書き i += 1; } else { // 左の子を足すとだめ // 特に何もしなくていいね } } return i - n;
一般の場合（左端を固定）
セグ木を完全二分木に限定せず、二分探索の始点を特に限定しない場合を考えます。始点を \(l\) とします。
セグ木で \([l, n)\) の区間和を求める際に見るノードを、左から順に見ていきます。 このノードたちの添字は、非再帰セグ木のループを使うと簡単に求められます（順序には注意）。
必要に応じて自分で図を書いたり、上のスライドの図を見るといいかもしれません。
これらのノードを順に見ていきながら、「ここまで加えると条件を満たさなくなる」という境界を探します。 このノードは \(O(\log(n))\) 個なので、単に線形探索するだけで問題ありません（ここすき）。
さて、そのようなノードが見つかったとします*2。 重要なこととして、あるノードの子たちを見ると、完全二分木になっています。 このことから、上で示した簡単な場合に帰着されるので、\(O(\log(n))\) 時間でできます。
以上を合わせて、\(O(\log(n))\) 時間で答えを求められます。
一般の場合（右端を固定）
上記で左から見ていたところを、右から見るようにするだけです。 off-by-one とかには注意。
遅延セグ木の場合
必要に応じて伝播しながらやると、同じ感じでできます。
おわり
候補が \(O(\log(n))\) 個なので線形探索をしても \(O(\log(n))\) 時間で済むパートすき。
