1日目: RadixHeapを実装する

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

RadixHeapって?

• 非負整数型のような型 (後述) である
• 最後に取り出した値より小さい値を入れることはできない

という条件がある代わりに速めな優先度付きキューです.

非負整数型のような型って?

上で書いた「非負整数型のような型」の要件を書きます.
全順序がある型$T$をRadixHeapに乗せるには次の条件を満たす関数$\def\dfunc{\operatorname{d}}\dfunc: T \times T \rightarrow \Bbb{N} \cup \set{0}$が存在して, ある程度高速に計算できる必要があります.

• $\dfunc$が返す値は小さめの上限$M$を持つ (計算量は$M$に比例します)
• $\forall x, y \in T [x = y \iff \dfunc(x, y) = 0]$
• $\forall x, y, z \in T [x < y < z \implies \dfunc(x, z) \le \dfunc(y, z)]$
• $\forall x, y, z \in T [\dfunc(x, y) < \dfunc(x, z) \implies \dfunc(x, z) = \dfunc(y, z)]$

非負整数型の場合は$\dfunc(x,y)=\lfloor\log_2((x\oplus y)+1)\rfloor$とすると条件を満たし, RustではT::BITS - (x ^ y).leading_zeros()のようにして計算することができます. このとき$M$の値はT::BITSになります.

仕組み

型TのRadixHeapは長さM + 1の「Tの配列」の配列buckets: [Vec<T>; M + 1]と最後に取り出した値last: Tを持ちます.
前述の通りlastよりも小さい値を新たにRadixHeapに入れることは出来ません.

要素をd(last, x)の値に応じてbucketで分類して管理するイメージです. buckets[i]の要素xについてd(last, x) == iが成り立ちます.

// サンプルコードではu32型の場合を書く
type T = u32;
fn d(x: T, y: T) {
    T::BITS - (x ^ y).leading_zeros()
}
// 関数dが返す値の最大値は今回は32
const M: usize = 32;

struct SampleRadixHeap {
    buckets: [Vec<T>; M + 1],
    last: T,
}

new (初期化)

bucketsには空の配列を詰めて, lastには型Tの最小値を入れておきます.

impl SampleRadixHeap {
    fn new() -> Self {
        Self {
            buckets: std::array::from_fn(|_| vec![]),
            last: T::MIN,
        }
    }
}

push

要素vをRadixHeapに入れるときはd(last, v)を計算し, その値でbucketsの中のどの配列に入れるか分類します.

impl SampleRadixHeap {
    fn push(&mut self, v: T) {
        assert!(v >= self.last);
        self.buckets[d(self.last, v)].push(v);
    }
}

pop

まず, bucketの中の配列で, 要素を持っている最初のものを探して取り出し, 代わりに空の配列を入れておきます.
取り出した配列はrakと呼ぶことにします

rak中から最小値を見つけてlastをその値で置き換え, そのlastの値に基づいてrakの要素をRadixHeapに入れ直します.

最後に, bucket[0]は必ず要素を持っているので, その値を取り出して返り値にします.

impl SampleRadixHeap {
    fn pop(&mut self) -> T {
        // 要素を持っている最初の配列を探して, 取り出す
        let rak = {
            let lak = self.buckets.iter_mut().find(|v| v.len() > 0).unwrap();
            std::mem::take(lak)
        };

        // 最小値を見つける
        self.last = *rak.iter().min().unwrap();

        // rakの要素を詰め直す
        for v in rak {
            self.buckets[d(self.last, v)].push(v);
        }

        // `buckets[0]`に要素が入ってるので, 取り出して返す
        return self.buckets[0].pop().unwrap();
    }
}

実装

• radixheap.rs

実装ではdをradix_distに, MをMAX_DISTと名前を変えています.
抽象化もばっちりしました.

まとめ

いかがでしたか?
明日はこのRadixHeapを使ってDijkstra法を実装します.