$$
\newcommand \lsb { \mathrm { lsb } }
$$
Fenwick tree (Binary indexed tree, BIT) をお勉強しましたところ、大変美しい構造を持つことがわかりましたから、共有させていただきたいと思います。案外長くなってしまって、自分でもびっくりしています。
第一章：Fenwick tree の基本機能
初めての方はまずこちらです。
問題
整数の配列 $a$ があります。これに次のようなクエリが飛んできます。
| 操作 | 内容 |
|---|---|
| $\mathrm { push } (x)$ | 後ろに $x$ を追加します。 |
| $\mathrm { prefix \_ sum } (i)$ | $\sum _ { j = 0 } ^ { i - 1 } a _ j$ を計算します。 |
| $\mathrm { add }(i, x)$ | $a _ i$ に $x$ を加算します。 |
Fenwick tree にできること
これらをそれぞれ、こちらの計算量で処理できるのが Fenwick tree なのです。
| 操作 | 計算量 |
|---|---|
| $\mathrm { push }$ | 実行時間は vector の再確保がありますから最悪 $\Theta ( N )$、加算の回数は最悪 $\lfloor \lg N \rfloor$ 回、償却定数回です。 |
| $\mathrm { prefix \_ sum }$ | 最悪 $\lfloor \lg N \rfloor$ 回、平均でも、$\lg N / 2$ 回程度の加算をします。 |
| $\mathrm { add }$ | 最悪 $\lfloor \lg N \rfloor$ 回、平均でも、$\lg N / 2$ 回程度の加算をします。 |
Fenwick tree の内部構造
前提知識：最下位ビット
非負整数 $i$ に対して、$\lsb (i)$ を次のように定めましょう。これは「最下位ビット（least significant bit, LSB)」と呼ばれており、引数を割り切る最大の 2 冪（引数が $0$ のときには $0$ ）のことを指します。
$$
\lsb ( i )=
\begin {cases}
0 & i = 0, \\
\max
\left \{
j \in \{ 1, 2, 4, \dots \}
\middle \vert
i \in j \cdot \mathbb N
\right \}
& i \neq 0
\end{cases}
$$
2 の補数表現を持つ符号付き整数型のビット演算をつかうと、i & -i のように書くことができます。
#[inline]
fn lsb(i: usize) -> usize {
let i = i as isize;
(i & -i) as usize
}
もしくは（$i$ が $1$ 以上のとき） -i と i - 1 がビット反転の関係にあることから、このようにもかけます。これならば符号付きを経由する必要がなくて、簡単です。（C++ では ~ ですが、Rust では ! がビット反転となりますので注意です！）
#[inline]
fn lsb(i: usize) -> usize {
if i == 0 {
0
} else {
i & !(i - 1)
}
}
また、@try_ru_cpp さんにコメントを頂いたのですが、このように書くのが Rust ではベストかもしれません。
#[inline]
fn lsb(i: usize) -> usize {
i & i.wrapping_neg()
}
また（$i$ が $1$ 以上のとき、） $i$ の LSB を下ろす操作は i &= i - 1 ともかけます。
Fenwick tree の内部構造
長さ $n$ の配列 $a$ に対して、長さ $n + 1$ の配列 $b = f ( a )$ を 次のように定義しましょう。これは累積和の開始位置を $0$ ではなく $i - \lsb (i)$ にしたものと思うとよいです。
$$
b _ i = \displaystyle \sum _ { j = 1 } ^ { \lsb ( i ) } a _ { i - j }
$$
ソースコード
new は空配列に対応する Fenwick tree （中身は $0$ がひとつだけある vector です。）を構築します。こちらの構造体に機能を追加していきましょう！
struct Fenwick {
table: Vec<i32>,
}
impl Fenwick {
pub fn new() -> Self {
Self { table: vec![0] }
}
}
クエリ実装編 〜 push
$x = a _ { n - 1 }$ と $b _ 0, \dots, b _ { n - 1 }$ から $b _ n$ を作れればヨシ！です。
数学パート
$n = ( 2 k + 1 ) \cdot 2 ^ p \ ( k, p \in \mathbb N )$ のように一意的に表すと、$n$ の LSB は $2 ^ p$ ですから、このようになります。
$$
b _ n
= \sum _ { i = 1 } ^ { 2 ^ p } a _ { n - i }
= x + \sum _ { q = 0 } ^ { p - 1 } \left ( \sum _ { i = 1 } ^ { 2 ^ q } a _ { n - 2 ^ q - i } \right )
= x + \sum _ { q = 0 } ^ { p - 1 } b _ { n - 2 ^ q }
$$
ソースコード
つまり、$d$ が $\lsb ( n )$ と異なる限り $\{ 1, 2, 4, \dots \}$ を走査していって、$b _ { n - d }$ を順に足していくと良いです。
pub fn push(&mut self, mut x: i32) {
let n = self.table.len();
let mut d = 1;
let k = lsb(n);
while d != k {
x += self.table[n - d];
d *= 2;
}
self.table.push(x);
}
結果確認のコツ
$1, 10, 100, 1000, \dots$ を push して中身を改行区切りで出力するとわかりやすいです。
0
1
11
100
1111
10000
110000
1000000
クエリ実装編 〜 prefix_sum
$\sum _ { j = 0 } ^ i a _ j$ を、$b$ の項をなんとかかき集めて作っていきましょう。式で書くとゴツゴツのアハンになるタイプのお話ですから、自然言語中心でご説明していこうと思います。
数学パート
まず $b _ i$ をつかうとこのように、長さ $\lsb ( i )$ だけ小さいものに帰着します。
$$
\sum _ { j = 0 } ^ { i - 1 } a _ j = b _ i + \sum _ { j = 0 } ^ { i - \lsb ( i ) } a _ j
$$
これを繰り返していくとどんどん小さくなり、最後には $0$ に帰着して終了となります。
計算量
「LSB 分だけ減らす」というのは結局、「LSB を下ろす」ということと同じですね。ということは、ビットの個数分以下であることがわかります。
ソースコード
つまりこのように $i$ から $\lsb ( i )$ を減らしていくと良いです。
pub fn prefix_sum(&self, mut i: usize) -> i32 {
let mut res = 0;
while i != 0 {
res += self.table[i];
i -= lsb(i);
}
res
}
もしくは便利なイテレータをつかうもよしです。
pub fn prefix_sum(&self, i: usize) -> u32 {
iter::successors(Some(i), |&i| Some(i - lsb(i)))
.take_while(|&i| i != 0)
.map(|i| self.table[i])
.sum()
}
クエリ実装編 〜 add
花形でしょうか。$b$ の項のうち $a _ i$ が寄与しているところを列挙です！ 実はこれがけっこう難しかったりします。
数学パート 1
そもそもいくつあるのでしょう。数えてみましょう！ $a _ i$ が $b _ j$ に寄与している条件は、$j = (2 k + 1 ) 2 ^ p \ ( k, p \in \mathbb N)$ と書いたときに
$$
2 k \cdot 2 ^ p \le i \lt ( 2 k + 1 ) \cdot 2 ^ p
$$
となることでした。おわかりですね。$p$ を固定するとこれを満たす $k$ は高々ひとつです。また、$n$ 以下の部分以外は管理しておらずですから、結局 $p \le \lg n$ なるもののみを見ると良いです。
突然ですが、大人召喚です。
大人の事情により $i$ に予め $1$ を足しておきますね。するとこのように、不等号についている等号が入れ替わる感じになります。
$$
2 k \cdot 2 ^ p \lt i \le ( 2 k + 1 ) \cdot 2 ^ p
$$
数学パート 2
まずこのような $k$ の存在する $p$ なのですが、式をよく見ると $\lfloor i / { 2 ^ p } \rfloor$ が偶数であることと同値であることがわかります。言い換えると、そのビットが降りていることと同値です。
それぞれの $p$ に対して、加算したい箇所は $( 2 k + 1 ) \cdot 2 ^ p$ でした。不等式をみると、これは $i$ 以上である最小の $2 ^ p$ の倍数であることがわかります。
数学パート 3
ところで、先程の不等式を満たす $i, k, p$ を考えたとき、$i$ と $(2 k + 1 ) \cdot 2 ^ p$ では $2 ^ p$ 以下のビットしか異ならないことに注意しましょう。言い換えると、小さい順にみて $(p, k)$ を見つけるたびに毎回 $i$ に $(2 k + 1) \cdot 2 ^ p$ を代入する操作をしてもあとの結果に影響しない ことがわかります。
この代入操作をすると何が嬉しいかというと、ある $p$ に関して $k$ を探しているときに、$i$ が 常に $2 ^ p$ の倍数になります。 その場合、$i = 2 k \cdot 2 ^ p$ ですから、目的の $(2 k + 1) \cdot 2 ^ p$ に到達するためには LSB を足せば良いことがわかります。
ソースコード
$i$ さんには、はじめに大人の下駄を履いていただいて、それ以降は LSB を足していくスイスイエスカレータコースで、第二宇宙速度を超えてお外に飛んでいくまで頑張っていただきましょう！
pub fn add(&mut self, mut i: usize, x: i32) {
i += 1;
while i < self.table.len() {
self.table[i] += x;
i += lsb(i);
}
}
もしくは便利なイテレータをつかうもよしです。
pub fn add(&mut self, i: usize, x: u32) {
let n = self.table.len();
iter::successors(Some(i + 1), |i| Some(*i + lsb(*i)))
.take_while(|&i| i < n)
.for_each(|i| self.table[i] += x);
}
構築
数学パート
push を繰り返す方法もあるのですが、それよりも書きやすい方法があります。$b _ i$ が $b _ j$ の push の巻き添えになる条件が $j = i + \lsb ( i )$ ですから、これを配る DP するとよいです。
ソースコード
このように、一重ループになり、コードも簡単です。
pub fn build(src: &[i32]) -> Self {
let mut table = vec![0; src.len() + 1];
for i in 1..table.len() {
let x = src[i - 1];
table[i] += x;
let j = i + lsb(i);
if j < table.len() {
table[j] += table[i];
}
}
Self { table }
}
イテレータを使うも良しです。
pub fn from_slice(src: &[u32]) -> Self {
let mut table = vec![0; src.len() + 1];
table[1..].copy_from_slice(src);
let n = table.len();
(1..n)
.map(|i| (i, i + lsb(i)))
.filter(|&(_, j)| j < n)
.for_each(|(i, j)| table[j] += table[i]);
Self { table }
}
第二章：グラフで理解する Fenwick tree
長さ $n$ の数列 $a$ に対して、数列 $b, c$ を
$$
\begin {gathered}
b _ i = \sum _ { j = 1 } ^ { \lsb ( i ) } a _ { i - j }
&
c _ i = \sum _ { j = 1 } ^ { i } a _ { i - j }
\end {gathered}
$$
（ただし $i = 0, \dots , n$）で定義しましょう。
つまり $b$ は Fenwick tree の内部配列で、$c$ は累積和です。
Fenwick tree graphs
$i$ から $i - \lsb ( i )$ に辺を張ったグラフと、$i + \lsb ( i )$ に辺を張ったグラフを考えましょう。縦軸が $i$、横軸が $\mathrm { trailing \_ zeros } ( i )$ になるように描画するときれいです。
左側のグラフと右側のグラフの関係
最大頂点が $2$ 冪になるように打ち切ると、互いに上下反転の関係になります。
また $\lsb ( -i ) = \lsb ( i )$ と思うことにすると、頂点番号を $-1$ 倍して同様の定義をすると、左のグラフは右のグラフに、右のグラフは左のグラフになります。
これらのグラフと数列 a, b, c の関係
左のグラフに関しては、各頂点 $i$ に $c _ i$ のポテンシャルをもたせましょう。右のグラフに関して、各頂点 $i$ に $a _ { i - 1 }$ の重みをもたせましょう。（off-by-one に注意です！）
すると、このように対応します。
| 数列 | 左側のグラフ | 右側のグラフ |
|---|---|---|
| $a _ i$ | $i$ と $i + 1$ のポテンシャルの差 | $i + 1$ の重み |
| $b _ i$ | $i$ から出ていく辺の両端のポテンシャルの差 | $i$ の部分木の重み |
| $c _ i$ | $0$ と $i$ のポテンシャルの差 | $i$ やそれより小さい添字の部分の重みの和 |
一方、各種操作は次のように言い換えることができます。
| 操作 | 内容 |
|---|---|
| $\mathrm { push } ( \_ ) $ | $b$ に挿入する値を計算します。 |
| $\mathrm { add } ( i, \_ ) $ | $a _ i$ を編集します。もしくは $c$ の suffix を編集すると見ても良いです。 |
| $\mathrm { prefix \_ sum } ( i, \_ ) $ | $c _ i$ を計算します。もしくは $a$ の prefix を足し合わせると見ても良いです。 |
しかし実際に直接編集できるのは $b$ だけですから、$a$ に関する操作や $c$ に関する操作を頑張って $b$ に関する操作に言い換える必要があります。
グラフで見る push
$\mathrm { push }$ は、$b _ i$ を $a _ { i - 1 }$ と愉快な仲間たちを使って計算します。
まずはこちらをご覧いただけるとです。$x = a _ { 23 }$ を挿入するために $b _ { 24 }$ を計算しようとしている図です。（アッこれ右側の図 push $a _ { 23 }$ の誤りですね。ごめんなさい、うっかりしておりました。）
左側のグラフ
計算したいものは青い点線の部分のポテンシャルの差ですから、別経路で計算すると良いです。
右側のグラフ
$b _ i$ は緑頂点の部分木の重みですから、子の部分木の重みを使うと計算できます。
push の償却計算量
方法 1 （左側のグラフ）
頂点 $0$ から最大頂点へのパスをスタックで管理すると思えばわかりやすいです。入る回数と出る回数は、頂点ごとにたかだか $1$ 回ずつです。
方法 2 （左側のグラフ）
もしくは、添字の隣接する頂点を（この図の緑辺）にすべて辺（二重辺になる場合もあります）を張っておくと、一度の push 操作が一つの領域に対応することがわかります（図割愛） から、各辺が隣接する領域が当然たかだか 2 つであることから、こちらも簡単に理解することができます。
方法 3 （右側のグラフ）
push するたびに連結成分を一つ追加して、巻き添えにした分だけ減っていきますから、結局巻き添えにした回数も高々 $n$ 回（より正確に言うと、$n$ から最終的な連結成分の個数を引いたもの）であることがわかります。
グラフで見る add
$\mathrm { add }$ は $a$ の一点に作用します。$c$ の suffix に作用すると考えます。
こちらの図をご覧ください。$a _ { 21 }$ に $x$ を足そうとしています。もしくは、$c$ の suffix に $x$ を足すと見ることもできます。
左側のグラフ
$22$ 以降のポテンシャルを足せばよいですから、$21$ と $22$ の間を跨いでいる辺全てに関して、$b$ に足し算をすると良いです。
右側のグラフ
$21$ の重みに足したいわけですから、先祖方向に向かって $b$ に足し算をすると良いです。
グラフで見る prefix_sum
$\mathrm { prefix\_sum }$ は $c _ i$ の一点を取得します。もしくは $a$ の prefix に関して足し合わせます。
こちらの図をご覧ください。$c _ { 28 }$ を取得しようとしています。もしくは $a$ の prefix の総和を取得しようとしています。
左側のグラフ
$0$ と $28$ とのポテンシャル差を計算したいのですから、先祖方向に向かって $b$ を足し合わせると良いです。
右側のグラフ
$28$ 以下の頂点を部分木に分解してみると、結局 $28$ と $29$ の間を跨いでる辺全てに関して、$b$ を足し合わせると良いです。
add と prefix_sum の計算量
先祖クライミングをみると、結局 $\lfloor { \lg n } \rfloor - \mathrm { trailing \_ zeros } ( i )$ 回の操作になることがわかります。
第三章：二分探索
いわゆる標準的な Fenwick tree の実装はここまでなのですが、Fenwick tree は実は二分探索もできてしまうのです。$a$ は非負であるなどとして、累積和が単調になるように仮定を増やしておきましょう。すると二分探索ができます。
定義
累積和上の二分探索の添字をどうするかについては、様々な流派がありそうです。（たとえば、こちらの定義に $1$ を足したものとするのにも妥当性がありそうです。）皆様のご意見を伺いたいところですが、とりあえずこちらの定義で行きましょう。
| 操作 | 要件 | 内容 |
|---|---|---|
| $\mathrm { upper \_ bound } (x)$ | $0 \le x$ | $\mathrm { prefix \_ sum } ( i ) \le x$ となる最大の $i \in \left \{ 0, \dots, \mathrm { len } ( b ) \right \}$ を返します。 |
やり方
引き算は必要ありません。結論から申し上げると、これの左側のグラフでそのまま（子を大きい順に線形探索する感じで）探索すればよいです。
動きをぐっと睨むと、普通の二分探索と全く同じになっていることがわかります。ダブリングで二分探索をするテクの $0$ 始点特化版だと思ってくださると嬉しいです。
ソースコード
変化量 $d$ を $2$ で割っていきながら $x$ 以下という条件を保ちながらなるべく大きくしていく感じです。
pub fn upper_bound(&self, x: &u32) -> usize {
let mut d = self.table.len().next_power_of_two() / 2;
let mut j = 0;
let mut now = 0;
while d != 0 {
if d + j < self.table.len() {
let next = now + self.table[d + j];
if &next <= x {
now = next;
j += d;
}
}
d /= 2;
}
j
}
prefix_sum と二分探索
prefix_sum も、$i$ を超えないという条件を保ちながらなるべく大きくしていく という方針（先程ご説明した add の実装とは、逆）をとることができます。それを踏まえると実は $\mathrm { prefix \_ sum }$ も類似の実装ができます。（余計ややこしいですが……）
pub fn prefix_sum(&self, i: usize) -> u32 {
let mut d = self.table.len().next_power_of_two() / 2;
let mut j = 0;
let mut now = 0;
while d != 0 {
if d + j <= i {
now += self.table[d + j];
j += d;
}
d /= 2;
}
now
}
というわけで、このような謎共通化ができます（？）
#[inline]
pub fn upper_bound(&self, x: &u32) -> usize {
self.partition_point(|_, y| y <= x).0
}
#[inline]
pub fn prefix_sum(&self, i: usize) -> u32 {
self.partition_point(|j, _| j <= i).1
}
pub fn partition_point(&self, pred: impl Fn(usize, &u32) -> bool) -> (usize, u32) {
let mut d = self.table.len().next_power_of_two() / 2;
let mut i = 0;
let mut now = 0;
while d != 0 {
if d + i < self.table.len() {
let next = now + self.table[d + i];
if pred(i + d, &next) {
now = next;
i += d;
}
}
d /= 2;
}
(i, now)
}
え、第三章もう終わりですか！？ そうなのです。二分探索が簡単なのが Fenwick tree のとても大きな強みのひとつです。
第四章：引き算で広がる世界
引き算があると、これが作れます！
| 操作 | 内容 |
|---|---|
| $\mathrm { sum } ( l, r )$ | $\sum _ { i = l } ^ { r - 1 } a _ i$ を計算します。 |
| $\mathrm { access } ( i )$ | $a _ i$ を取得します。 |
| $\mathrm { set } ( i, x )$ | $a _ i$ を $x$ に変更します。 |
| $\mathrm { partition \_ point } (l, \mathrm { pred })$ | $\mathrm { pred } ( \mathrm { sum } ( l, r ) )$ となる最大の $r \in \{ 0, \dots \mathrm { len } ( b ) \}$ を返します。 |
大して広がらないことがおわかりいただけたでしょうか。（炎上）
$\mathrm { sum }$ は $\mathrm { prefix \_ sum }$ の差を取るだけですし、$\mathrm { access } ( i )$ は $\mathrm { sum } ( i + 1, i )$ で、$\mathrm { set } ( i, x )$ は $\mathrm { add } ( i, x - \mathrm { get } ( i ) )$、$\mathrm { partition \_ point }$ は $\mathrm { pred } ( c _ i - c _ j)$ をするだけでしょう。
もちろん $\mathrm { upper \_ bound }$ も同じように作れるのですが、これに関してはそもそも $c _ i - c _ j \le x$ を移項して $c _ i \le c _ j + x$ とすれば引き算は必要ありません。（もちろん、定数を足す操作が順序を保つならばのお話です。）
第五章：二次元 Fenwick tree
長方形型二重配列
$h, w$ を非負整数としましょう。$i = 0, \dots, h - 1, \ j = 0, \dots, w - 1$ に対して値 $a _ { i , j }$ が定まっているとき、$a$ は高さ $h$、幅 $w$ の（長方形型）二重配列（二重数列）であるといいます。この章では単に二重配列といえば長方形型のことを指すものとします。
問題
二重配列 $a$ に対して、次のような操作をしたいです。
| 操作 | 内容 |
|---|---|
| $\mathrm { double \_ prefix \_ sum } (i, j)$ | $\sum _ { k = 0 } ^ { i - 1 } \sum _ { l = 0 } ^ { j - 1 } a _ { k , l }$ を計算します。 |
| $\mathrm { add }(i, j, x)$ | $a _ { i, j }$ に $x$ を加算します。 |
二次元 Fenwick tree にできること
こちらが計算量となっております。
| 操作 | 計算量 |
|---|---|
| $\mathrm { double \_ prefix \_ sum }$ | 最悪 ${ \lfloor \lg N \rfloor } ^ 2$ 回、平均 $( \lg N / 2 ) ^ 2$ 回程度の加算が必要です。 |
| $\mathrm { add }$ | 最悪 ${ \lfloor \lg N \rfloor } ^ 2$ 回、平均 $( \lg N / 2 ) ^ 2$ 回程度の加算が必要です。 |
二次元 Fenwick tree の内部構造
こうです。
$$
b _ { i , j } = \sum _ { k = 1 } ^ { \lsb ( i ) } \sum _ { l = 1 } ^ { \lsb ( j ) } a _ { k , l }
$$
中身は vector の vector となります。
pub struct Fenwick2d {
pub table: Vec<Vec<u32>>,
}
impl Fenwick2d {
pub fn new() -> Self {
Self {
table: vec![vec![0]],
}
}
}
クエリ実装編 〜 double_prefix_sum
普通に二重ループをするとよいです。
pub fn double_prefix_sum(&self, mut i: usize, j: usize) -> u32 {
let mut res = 0;
while i != 0 {
let mut j = j;
while j != 0 {
res += self.table[i][j];
j -= lsb(j);
}
i -= lsb(i);
}
res
}
イテレータを使うとこのように書くこともできます。
pub fn double_prefix_sum(&self, i: usize, j: usize) -> u32 {
itertools::cartesian_product(
iter::successors(Some(i), |&i| Some(i - lsb(i))).take_while(|&i| i != 0),
iter::successors(Some(j), |&j| Some(j - lsb(j))).take_while(|&j| j != 0),
)
.map(|(i, j)| self.table[i][j])
.sum()
}
クエリ実装編 〜 add
こちらも普通に二重ループを回すだけです。（突然変な関数を使っていますが、h, w は $b$ の縦と横の大きさだと思っていただけるとです。）
pub fn add(&mut self, mut i: usize, mut j: usize, x: u32) {
let (h, w) = exact_size_of_grid(&self.table);
i += 1;
j += 1;
while i < h {
let mut j = j;
while j < w {
self.table[i][j] += x;
j += lsb(j);
}
i += lsb(i);
}
}
イテレータを使うとこのように書くこともできます。
pub fn add(&mut self, i: usize, j: usize, x: u32) {
let (h, w) = exact_size_of_grid(&self.table);
itertools::cartesian_product(
iter::successors(Some(i + 1), |&i| Some(i + lsb(i))).take_while(|&i| i < h),
iter::successors(Some(j + 1), |&j| Some(j + lsb(j))).take_while(|&j| j < w),
)
.for_each(|(i, j)| self.table[i][j] += x);
}
構築
累積和と同じで、やり方が二通りあります。
- 方法 1: 一旦全箇所の $i$ 方向の遷移をしてから $j$ 方向の遷移をします。（配る DP）
- 方法 2: 一箇所ずつ、$i$ 方向と $j$ 方向それぞれに遷移して、超く部分を除きます。（貰う DP）
どちらにもどちらの代えがたいメリットがあります。方法 1 は引き算が必要ありません。方法 2 は push に応用ができます。
構築方法 1
かなり長いのですが、このように 1 方向ずつ遷移していくと重複がでなくてよいです。
pub fn from_slice_vec(src: &[Vec<u32>]) -> Self {
let (h, w) = exact_size_of_grid(src);
let mut table = vec![vec![0; w + 1]; h + 1];
for i in 1..=h {
let next_i = i + lsb(i);
if next_i <= h {
for j in 1..=w {
let x = table[i][j];
table[next_i][j] += x;
}
}
}
for j in 1..=w {
let next_j = j + lsb(j);
if next_j <= w {
for i in 1..=h {
table[i][j] += src[i - 1][j - 1];
let x = table[i][j];
table[i][next_j] += x;
}
}
}
Self { table }
}
イテレータを使うとこのようになります。余計分かりづらくなったのは内緒です。
pub fn from_slice_vec(src: &[Vec<u32>]) -> Self {
let (h, w) = exact_size_of_grid(src);
let mut table = iter::once(vec![0; w + 1])
.chain(
src.iter()
.map(|v| iter::once(0).chain(v.iter().copied()).collect::<Vec<_>>()),
)
.collect::<Vec<_>>();
let mut table = vec![vec![0; w + 1]; h + 1];
for ((i, next_i), j) in itertools::cartesian_product(
(1..=h)
.map(|i| (i, i + lsb(i)))
.filter(|&(_, next_i)| next_i <= h),
1..=w,
) {
let x = table[i][j];
table[next_i][j] += x
}
for ((j, next_j), i) in itertools::cartesian_product(
(1..=w)
.map(|j| (j, j + lsb(j)))
.filter(|&(_, next_j)| next_j <= w),
1..=h,
) {
let x = table[i][j];
table[i][next_j] += x
}
Self { table }
}
構築方法 2
ジャギー Fenwick tree の push でお見せいたします。$\mathrm { push ( i, x ) }$ は $i ( \le \mathrm { len } ( a ) )$ 列目に $x$ を挿入します。
pub fn push(&mut self, i: usize, mut x: u32) {
assert!(i < self.table.len());
let i = i + 1;
if i == 1 {
self.table[0].push(0);
}
if i == self.table.len() {
self.table.push(vec![0]);
}
let j = self.table[i].len();
iter::successors(Some(lsb(i)), |&d| Some(d / 2)).for_each(|d| x += self.table[i - d][j]);
iter::successors(Some(lsb(j)), |&e| Some(e / 2)).for_each(|e| x += self.table[i][j - e]);
itertools::cartesian_product(
iter::successors(Some(lsb(i)), |&d| Some(d / 2)),
iter::successors(Some(lsb(j)), |&e| Some(e / 2)),
)
.for_each(|(d, e)| x -= self.table[i - d][j - e]);
self.table[i].push(x);
}
二分探索
いろいろなパターンがありそうですが、今回は
$$
\mathrm { horizontal \_ upper \_ sum } ( i, x )
= \left \{
j \in \{ 0, \dots, \mathrm { len } ( a ) \}
\middle \vert
\mathrm { double \_ prefix \_ sum } ( i, j ) \le x
\right \}
$$
を計算することにしましょう。
pub fn horizontal_upper_bound(&self, i: usize, x: &u32) -> usize {
let table_width = exact_size_of_grid(&self.table).1;
let mut j = 0;
let mut now = 0;
for d in iter::successors(Some(table_width.next_power_of_two() / 2), |&d| Some(d / 2))
.take_while(|&d| d != 0)
{
if j + d < table_width {
let next = now + self.i_prefix_sum_j_raw_element(i, j + d);
if &next <= x {
j += d;
now = next;
}
}
}
j
}
fn i_prefix_sum_j_raw_element(&self, i: usize, j: usize) -> u32 {
iter::successors(Some(i), |&i| Some(i - lsb(i)))
.take_while(|&j| j != 0)
.map(|i| self.table[i][j])
.sum()
}
引き算で広がる世界 − 外伝
2 次元 Fenwick tree の二分探索で、$i$ の始点を $0$ に限らない感じにできます。
あとがき
いかがでしたでしょうか！
ご助言くださったツイッターの皆様です。ありがとうございました。