セグメント木 (src/data_structure/segment_tree.rs)

View this file on GitHub
View document part on GitHub
Last update: 2024-04-26 00:41:30+09:00

`segment_tree`

`SegmentTree`

数列を管理しながら1点更新・区間演算を高速に行うデータ構造．

前提条件

集合 $S$ がある演算 $\cdot : S\times S\to S$ に関してモノイドをなすとする．すなわち

（結合律）任意の $x, y, z\in S$ に対して $(xy)z=x(yz)$
（単位元の存在）ある $e\in S$ が存在して任意の $x\in S$ に対して $ex=xe=x$

が成り立つと仮定する．

このとき，$S$ の元からなる列 $a_0,\cdots a_{N-1}\; (a_i\in S)$ に対して

要素の1点更新
区間の要素の総積の取得

をいずれも $O(\log N)$ で行うことができる．

`new`

fn new(n: usize, id: T, op: F) -> Self

n 個の要素からなる数列を構築する．
id: T は $S$ の単位元
op: F は $S$ に対して定義された二項演算
- F は Fn(T, T) -> T トレイトを実装している．
初期値はすべて id である．

制約

$n\geq 0$

計算量

$O(n)$

`construct`

fn construct(&mut self, seq: &Vec<T>)

数列の値を seq の要素で置き換える．

制約

seq.len()は数列の長さ $n$ に等しい

計算量

$O(n)$

`update`

fn update(&mut self, mut i: usize, x: T)

数列の i 番目の値を x で置き換える．

制約

$0\leq i < n$

計算量

$O(\log n)$

`prod`

fn prod(&self, mut l: usize, mut r: usize) -> T

区間 [l, r) のすべての要素にわたる積を返す．

制約

$0\leq l \leq r \leq n$

計算量

$O(\log n)$

`get`

fn get(&self, i: usize) -> T

数列の i 番目の値を返す．

制約

$0 \leq i < n$

計算量

$O(1)$

Depends on

Required by

Verified with

Code

pub mod segment_tree {
    pub struct SegmentTree<T, F> {
        id: T,
        op: F,
        size: usize,
        values: Vec<T>,
    }

    #[allow(dead_code)]
    impl<T: Copy, F: Fn(T, T) -> T> SegmentTree<T, F> {
        pub fn new(n: usize, id: T, op: F) -> Self {
            //! Constructs a segment tree `a[i]` for array size `n`.
            SegmentTree {
                id: id,
                op: op,
                size: n,
                values: vec![id; 2 * n],
            }
        }

        pub fn construct(&mut self, seq: &Vec<T>) {
            //! Constructs a segment tree `a[i]` from a vector `seq`.
            assert!(seq.len() == self.size);

            for (i, &x) in seq.iter().enumerate() {
                self.values[i + self.size] = x;
            }

            for i in (1..self.size).rev() {
                self.values[i] = (self.op)(self.values[i << 1], self.values[i << 1 | 1]);
            }
        }

        pub fn update(&mut self, mut i: usize, x: T) {
            //! Makes a point update: `a[i] <- x`.
            assert!(i < self.size);

            i += self.size;
            self.values[i] = x;
            while i > 1 {
                i >>= 1;
                self.values[i] = (self.op)(self.values[i << 1], self.values[i << 1 | 1]);
            }
        }

        pub fn prod(&self, mut l: usize, mut r: usize) -> T {
            //! Returns a product over the interval `[l, r)`.
            assert!(l <= r && r <= self.size);

            l += self.size;
            r += self.size;

            let mut prod_l = self.id;
            let mut prod_r = self.id;

            while l < r {
                if l & 1 == 1 {
                    prod_l = (self.op)(prod_l, self.values[l]);
                    l += 1;
                }
                if r & 1 == 1 {
                    r -= 1;
                    prod_r = (self.op)(self.values[r], prod_r);
                }

                l >>= 1;
                r >>= 1;
            }

            (self.op)(prod_l, prod_r)
        }

        pub fn get(&self, i: usize) -> T {
            //! Returns a value of `a[i]`.
            assert!(i < self.size);

            self.values[i + self.size]
        }
    }
}