18日目: Persistent SegmentTreeを実装する

競プロ用ライブラリを作る Advent Calendar 2024の17日目です.

Persistent SegmentTreeって?

過去の情報を取得できるセグ木です.
ここでは「コピーのコストが低いセグ木」として実装します.

仕組み

SegmentTreeは二分木のような構造で列を管理するデータ構造でした.
そしてPersistentArrayは木構造で列を管理して複製コストを抑えるデータ構造でした.

ということで, この2つを組み合わせます. 上手く書けば通常のSegmentTreeよりも実装が楽です.

// サンプルコードではi32の足し算を計算する永続セグ木を実装します
#[derive(Clone)]
enum SamplePersistentSegTree {
    // 末端要素
    Leaf(i32),
    // 要素数が2以上の列. 永続配列と同じように, 2つの列の表現のペア + 付加情報で表現する
    Trunc {
        // 要素数
        len: usize,
        // 左の部分木
        left: Rc<Self>,
        // 右の部分木
        right: Rc<Self>,
        // 全体での合計
        sum: i32,
    },
}
use SamplePersistentSegTree::*;

impl SamplePersistentSegTree {
    // 要素数を返す
    fn len(&self) -> usize {
        match self {
            // 要素が1つなので, 1を返す
            Leaf(_) => 1,
            // 要素数を保存してるので, それを返す
            Trunc { len, .. } => *len,
        }
    }

    // 全体での合計を返す
    fn all_sum(&self) -> &i32 {
        match self {
            // 要素が1つだけなので, その値をそのまま返す
            Leaf(v) => v,
            // 合計の値を保存しているので, それを返す
            Trunc { sum, .. } => sum,
        }
    }

    // indexに対応する要素を書き換える
    fn set(&mut self, index: usize, item: i32) {
        assert!(index < self.len());
        match self {
            // 末端の要素に辿り着いたら値を書き換える
            Leaf(v) => *v = item,
            Trunc {
                left, right, sum, ..
            } => {
                let left_len = left.len();
                // 目的の値が左右どちらに含まれてるか調べて, 適切な方に書き換えクエリを送る
                if index < left_len {
                    Rc::make_mut(left).set(index, item);
                } else {
                    Rc::make_mut(right).set(index - left_len, item);
                }
                // 全体の合計値を再計算する
                *sum = left.all_sum() + right.all_sum();
            }
        }
    }

    // 区間[l, r)での合計を返す
    fn range_sum(&self, l: usize, r: usize) -> i32 {
        assert!(l <= r && r <= self.len());
        // l == r なら区間に何も含まれないので0を返します
        if l == r {
            return 0;
        }
        // 全範囲の合計を求められていた場合はadd_sumを返します
        if l == 0 && r == self.len() {
            return *self.all_sum();
        }
        let Trunc { left, right, .. } = self else {
            // 1つしか要素がないRelayの場合は, 上の場合分けに必ず引っかかります
            unreachable!();
        };
        let left_len = left.len();
        // 左右の部分木と被っている区間の情報を取ってきて, その合計を返します
        left.range_sum(l.min(left_len), r.min(left_len))
            + right.range_sum(l.saturating_sub(left_len), r.saturating_sub(left_len))
    }
}

実装

• persistentsegmenttree.rs

上のサンプルプログラムを抽象化して, ちょっと無駄な部分を削りました.

まとめ

いかがでしたか?

明日は天才文字列アルゴリズムを実装します.