モノイド \((X, \times, 1)\) と \(X\) 上の数列
\[(x_1, \ldots, x_N \mid x_i \in X)\]について, 次の操作がそれぞれ \(O(\log N)\) で出来る.
区間とその積をノードに持つ (完全) 二分木を構成する. つまり, 区間 \(I\) に対応するノード \(\mathcal{N}_I\) があって, その積 \(p_I = \prod_{i \in I} x_i\) を 予め 計算しておく.
区間 \(I\) が空でないならば, それを左右に分割した区間 \(J_1\) と \(J_2\) に分割できるだろう（二等分が望ましい）. という訳で \(\mathcal{N}_I\) の子ノードとして \(\mathcal{N}_{J_1}\) と \(\mathcal{N}_{J_2}\) を持たせる. もちろん \(p_{J_1}, p_{J_2}\) も予め計算しておく.
根を全体区間 \([0, N)\) とし, 葉を一つの要素からなる区間 \(\{ x_i \}\) とする二分木が構成できたら, これがセグメントツリーである. セグメントツリーが持つノード数（つまり区間数）は \(2N\) 程度.
予め計算しておきたい \(p\) は再帰的に
\[p_I = p_{J_1} \times p_{J_2}\]とできるので, 葉から辿ってくことで, 積演算を \(2N\) 回すればよく, 木の構築は \(O(N)\) で出来る.
任意の区間積の計算は, 求めたい区間をちょうど被覆するノード \(O(\log N)\) 個程度に必ず分解できる. （これは \(N\) 程度の自然数は二進数で \(\log_2 N\) 桁で表せることに対応している.） 含むノードを大きいものから貪欲に選んでいけばよい.
更新は, 葉ノードを更新した後に, その親ノードを辿っていき \(p\) を更新していけばよい. このとき辿るノード数は木の高さなので \(O(\log N)\) である.
コンストラクタとして, 長さだけ指定して単位元で初期化する new と, 数列を vec で与える from とがある.
/// Sequence - Segment Tree
use crate::algebra::monoid::*;
pub struct SegmentTree<X> {
length_upper: usize, // power of 2
size: usize, // of nodes
data: Vec<X>,
}
impl<X> std::ops::Index<usize> for SegmentTree<X> {
type Output = X;
fn index(&self, i: usize) -> &Self::Output {
&self.data[self.size / 2 + i]
}
}
impl<X: Copy + Monoid> SegmentTree<X> {
pub fn new(length: usize) -> Self {
let mut length_upper = 1;
while length_upper < length {
length_upper <<= 1
}
let size = length_upper * 2 - 1;
let data = vec![X::one(); size];
SegmentTree {
length_upper,
size,
data,
}
}
pub fn from(xs: Vec<X>) -> Self {
let mut tree = Self::new(xs.len());
for i in 0..xs.len() {
tree.data[tree.size / 2 + i] = xs[i];
}
for i in (0..tree.size / 2).rev() {
tree.data[i] = tree.data[2 * i + 1] * tree.data[2 * i + 2];
}
tree
}
pub fn to_vec(self) -> Vec<X> {
self.data[self.size / 2..].to_vec()
}
pub fn update(&mut self, i: usize, t: X) {
let mut u = self.size / 2 + i;
self.data[u] = t;
while u > 0 {
u = (u - 1) / 2;
self.data[u] = self.data[u * 2 + 1] * self.data[u * 2 + 2];
}
}
fn product_sub(
&self,
range: std::ops::Range<usize>,
u: usize,
focus: std::ops::Range<usize>,
) -> X {
if focus.end <= range.start || range.end <= focus.start {
X::one()
} else if range.start <= focus.start && focus.end <= range.end {
self.data[u]
} else {
let mid = (focus.start + focus.end) / 2;
let a = self.product_sub(range.clone(), u * 2 + 1, focus.start..mid);
let b = self.product_sub(range.clone(), u * 2 + 2, mid..focus.end);
a * b
}
}
pub fn product(&self, range: std::ops::Range<usize>) -> X {
self.product_sub(range, 0, 0..self.length_upper)
}
}
impl<X: std::fmt::Debug> SegmentTree<X> {
pub fn debug(&self) {
#[cfg(debug_assertions)]
for i in 0..self.size {
if i > 0 && (i + 1).count_ones() == 1 {
eprintln!();
}
eprint!("{:?} ", &self.data[i]);
}
eprintln!();
}
}