16日目: PersistentArrayを実装する

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

PersistentArrayって?

過去に入っていた値なども見たりできる配列です.
ここでは「複製が低コストで行える配列」として実装します.
コストですが, 要素アクセス(取得/変更)も複製も$O(\log N)$で行えます ($N$は配列の要素数)

仕組み

配列の複製にコストがかかる要因は「複製するには要素を全部コピーしないといけない」というところにあります.
なので, 値を良い感じに使いまわせる構造にすればいいです.

どうするかというと, 列を木で表現します. これは二分木ですが, 実際はメモリキャッシュの大きさ等に合わせて多分木にすると高速化になります.

これで, 例えば最後の要素を「x」に置き換えた新たな配列を作るにはこうします.

こういう風に, 使いまわせるところを使いまわしていきます.

実装するときは, RustのRcみたいな共有ポインタを使います (C++ならshared_ptrです).
配列の複製は根本の参照をコピーして, 要素の変更をするときに変更したい要素のところをコピーして書き換えるイメージです.

// サンプルはi32の永続配列を実装します
#[derive(Clone)]
enum SamplePArray {
    // 要素1つ
    Val(i32),
    // 2つの配列の組み合わせ
    Array {
        // 合わせた配列の長さ
        len: usize,
        // 左の部分木
        left: Rc<Self>,
        // 右の部分木
        right: Rc<Self>,
    },
}

impl SamplePArray {
    // 配列の長さを返す
    fn len(&self) -> usize {
        match self {
            SamplePArray::Val(_) => 1,
            SamplePArray::Array { len, .. } => *len,
        }
    }

    // 取得するクエリ
    fn get(&self, index: usize) -> &i32 {
        match self {
            // 要素が1つだからそれを返す
            SamplePArray::Val(v) => v,
            SamplePArray::Array { len, left, right } => {
                // どちらの部分木に要素が含まれるか判定して,
                // その部分木から要素を取ってくる
                if index < *len {
                    left.get(index)
                } else {
                    right.get(index - len)
                }
            }
        }
    }

    // 可変参照を取得するクエリ
    fn get_mut(&mut self, index: usize) -> &mut i32 {
        match self {
            SamplePArray::Val(v) => v,
            SamplePArray::Array { len, left, right } => {
                // Rc::make_mutで, 必要に応じてデータを複製します.
                // 他はgetと変わらないです
                if index < *len {
                    Rc::make_mut(left).get_mut(index)
                } else {
                    Rc::make_mut(right).get_mut(index)
                }
            }
        }
    }
}

実装

• persistentarray.rs

余裕があったので末尾追加/削除もできるようにしてみました. Rustだと綺麗に実装すると型が要件を満たしていることを保証してくれていいですね. 安心感があります.

まとめ

いかがでしたか?

明日は永続シリーズを引き続き実装します.