区間和クエリを確かめる場合、$1$で初期化してておくと、セグメントの和とサイズが一致する。

#![allow(unused)]
fn main() {
    #[test]
    fn ones() {
        const MAX_SIZE: isize = 200;
        const OFFSET: isize = 10;

        for size in 0..MAX_SIZE {
            let range_sum_query =
                SegmentTree::<Add<isize>>::from_iter(std::iter::repeat_n(1, size as usize));
            for start in 0..=size {
                for sum in -OFFSET..=size as isize + OFFSET {
                    assert_eq!(
                        range_sum_query.partition_end(start as usize, |&v| v <= sum),
                        (start + sum).clamp(start, size) as usize,
                        "size: {size}, start: {start}, sum: {sum}"
                    )
                }
            }
        }
    }
}

#![allow(unused)]
fn main() {
pub trait Monoid {
    /// The set of the monoid.
    type Set;

    /// If [`Self::combine`] is commutative, some operations can be optimized.
    ///
    /// If unsure about the commutativity, use [`false`] for safety.
    ///
    /// # Commutative low
    ///
    /// ```text
    /// a · b = b · a    ∀ a, b ∈ Set
    /// ```
    const IS_COMMUTATIVE: bool;

    /// Returns the identity element.
    fn identity() -> Self::Set;

    /// Combines the two elements and returns the result.
    ///
    /// # Warning
    ///
    /// If the operation is **not** commutative, the position of the arguments matters.
    fn combine(lhs_or_prev: &Self::Set, rhs_or_new: &Self::Set) -> Self::Set;
}
}

Setをジェネリック型ではなく関連型としてもっています。 こうすることで、ジェネリックパラメーターが1つ減り、実装がシンプルになります。

#![allow(unused)]
fn main() {
pub struct SegmentTree<Query>
where
    Query: Monoid,
{
    /// Use `Box<T>` because the length is significant as follows.
    ///
    /// - data\[0\]    : dummy node (meaningless)
    /// - data\[1..n\] : nodes to store the combined value of the children.
    /// - data\[n..2n\]: nodes to store value for each cell.
    data: Box<[<Query as Monoid>::Set]>,

    /// `len` (number of elements) and offset (dummy + cache)
    len_or_offset: usize,
}
}

#![allow(unused)]
fn main() {
pub trait MonoidAction {
    /// The set of the monoid for update operation.
    type Map: Monoid;
    /// The set of the monoid for query operation.
    type Set: Monoid;

    /// Set [`true`] to use the segment size in [`Self::act()`]
    const USE_SEGMENT_SIZE: bool;

    /// Acts the mapping on the element and returns the result.
    ///
    /// # Size dependency
    ///
    /// You can access segment size if you want.
    /// This is equivalent to attaching the segment size information to [`Self::Set`] but more performant.
    fn act(
        mapping: &<Self::Map as Monoid>::Set,
        element: &<Self::Set as Monoid>::Set,
        size: Option<usize>,
    ) -> <Self::Set as Monoid>::Set;
}
}

Mapは更新クエリ、Setは取得クエリに対応します。

#![allow(unused)]
fn main() {
pub struct LazySegmentTree<Action>
where
    Action: MonoidAction,
{
    data: Box<[<<Action as MonoidAction>::Set as Monoid>::Set]>,
    lazy: Box<[<<Action as MonoidAction>::Map as Monoid>::Set]>,

    /// calculate if [`MonoidAction::USE_SEGMENT_SIZE`] is `true`.
    segment_size: Option<Box<[usize]>>,
}
}

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Add<T>(PhantomData<T>);

impl<T> Monoid for Add<T>
where
    T: Zero,
    for<'a> &'a T: std::ops::Add<Output = T>,
{
    type Set = T;

    const IS_COMMUTATIVE: bool = true;

    fn identity() -> Self::Set {
        T::zero()
    }

    fn combine(lhs_or_prev: &Self::Set, rhs_or_new: &Self::Set) -> Self::Set {
        lhs_or_prev + rhs_or_new
    }
}
}

例えば、Add<i32>はモノイド(i32, +, 0)に対応します。

トレイト境界のとり方には他にもいくつかのパターンが考えられます。

• T: Copy + Zero + std::ops::Add<Output = T>
• T: Clone + Zero + std::ops::Add<Output = T>

これらは、combine()の挙動に影響を与えます。 最適な実装は型に依るため、ユーザーに定義してもらうことにしました。

#![allow(unused)]
fn main() {
pub struct DynamicSegmentTree<Query>
where
    Query: Monoid,
{
    arena: Vec<Node<<Query as Monoid>::Set>>,
    range: Range<isize>,

    // save allocation cost
    reusable_stack: Vec<usize>,
}
}

Node<T>の詳細は後述します。

非可換モノイドでは要素の合成順が重要です。 後述のように、正しい順序で要素を合成するためにスタックを使用します。 reusable_stackはスタックの長さが$\lceil \log N\rceil$で抑えられることを利用して、アロケーションコストを節約するためにあります。

#![allow(unused)]
fn main() {
    pub fn with_capacity(range: Range<isize>, q: usize) -> Option<Self> {
        if range.is_empty() {
            None
        } else {
            // never panic: `range.len()` is always larger than 0
            let height = range.len().ilog2() as usize + 1;
            Some(Self {
                arena: Vec::with_capacity(q * height),
                reusable_stack: Vec::with_capacity(height * 4),
                range,
            })
        }
    }
}

seg_libはcrates.ioで公開しているため、返り値にOption型を取ることにしました。 競プロなどではpanic!()するのもありです。

#![allow(unused)]
fn main() {
struct Node<T> {
    index: isize,
    element: T,
    /// may be `None` if `combined == element`, avoiding `clone()`
    combined: Option<T>,

    left_ptr: Option<NonZeroUsize>,
    right_ptr: Option<NonZeroUsize>,
}
}

• seg_libの動的セグメント木では区間幅を固定しているので、各ノードが持つ区間情報は暗黙的に与えられます。 区間情報は根から子を辿る際に計算できるので、ノードには持ちません。
• ノードに対応する区間の途中結果にはOption型を採用しました。 これはトレイト境界<Query as Monoid>::Set: Cloneを回避するためです。
• 通常Option型は追加のフラグに8byte（64bitシステムの場合）利用しますが、Option<NonZero<T>>では0をNoneと扱う最適化が施されており、追加のメモリは必要ありません。 arenaの0要素目は根なので、子のインデックスは常に正です。

#![allow(unused)]
fn main() {
        while let Some(ptr) = self.reusable_stack.pop() {
            const MSB: usize = 1_usize.rotate_right(1);
            res = if ptr & MSB == 0 {
                <Query as Monoid>::combine(self.arena[ptr].get_element(), &res)
            } else {
                <Query as Monoid>::combine(&res, self.arena[!ptr].get_element())
            }
        }
}

resはタイプ2の場合の内側の計算結果です。

左側のノードへのポインターはそのまま、右側のノードへのポインターはビット反転してから、reusable_stackに格納しています。