rustで優先度付きキューを使いたい時、BtreeMapも検討してみようというお話

通常、rustで優先度付きキューを使うならBinaryHeapだとgoogle先生もchatGPT先生も言っています。

しかし、特定の用途に関してはBtreeMapを使うことになると思います。

AtCoderで問題を解いていた時に、この用途に出会ったので目線がAtCoderよりになっている可能性があります。

問題は↓です。

C - Cards Query Problem

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

【結論】優先度付きキューの中身を途中で見たいならBtreeMapを使う

【結論】優先度付きキューの中身を途中で見たいならBtreeMapを使う

ということで特定の用途とは、 優先度付きキューの中身を途中で見たいということです。

なぜなのか、理由を書いていきます。

BinaryHeapには優先度付きキューの中身を見るメソッドはない

BinaryHeapには、優先度付きキューの中身がどうなっているか見ることができません。

intosorted_vecを使うとキューが消費される

ソート済みのvecやiterに変換するメソッドはありますが、これは元のキューを消費してしまいます。例えば、以下のコードはコンパイルエラーになってしまいます。qをvecに変換した時、qは消費されてしまいなくなってしまいます。なので、新たに要素を追加するときコンパイルラーになってしまいます。

use std::collections::BinaryHeap;

fn main() {

    let mut q = BinaryHeap::new();
    q.push(1);
    q.push(2);
    q.push(3);

    let v = q.into_sorted_vec();
    for i in v.iter() {
        eprintln!("i = {:?}", i);
    }
    q.push(4);
}

エラーメッセージ

error[E0382]: borrow of moved value: `q`
   --> src/main.rs:14:5
    |
5   |     let mut q = BinaryHeap::new();
    |         ----- move occurs because `q` has type `BinaryHeap<i32>`, which does not implement the `Copy` trait
...
10  |     let v = q.into_sorted_vec();
    |               ----------------- `q` moved due to this method call
...
14  |     q.push(4);
    |     ^^^^^^^^^ value borrowed here after move

cloneして取っておけばいいじゃん、となりますがcloneを使うと実行速度が遅すぎて驚いてしまいますので、お勧めできません。

for文の中で使ったりするとTLEしてしまいます。

iterを使うとソートされずに出力される

iterを使うとソートされずに出力されてしまいます。

これは公式の例ですが、実行すると以下のようになります。

use std::collections::BinaryHeap;
let heap = BinaryHeap::from([1, 2, 3, 4]);

// Print 1, 2, 3, 4 in arbitrary order
for x in heap.iter() {
    println!("{x}");
}

出力結果

この辺りは予想になってしまいますが、BinaryHeapを実現するときのデータ格納の仕方に関係があると思っています。

BinaryHeapに関しては、参考になる記事がたくさんあると思いますので、ここでは割愛します。というか、BinaryHeapあまり理解してない…

BTreeMapを使って優先度付きキューを模倣する

模倣すると言っても、要は昇順(か降順)に値が取り出せればいいので、普通にMapに値を代入するだけです。

use std::collections::BTreeMap;

fn main() {
    let mut bt_map: BTreeMap<_, _> = BTreeMap::new();

    for i in [1, 2, 3, 4, 1] {
        bt_map.entry(i).and_modify(|e| *e += 1).or_insert(1);
    }

    for i in bt_map.iter() {
        eprintln!("i = {:?}", i);
    }
}

実行結果

i = (1, 2)
i = (2, 1)
i = (3, 1)
i = (4, 1)

注意するべき点としては、「あくまでもMapなのでkeyの重複ができない」ことです。なので、keyが重複しないという条件が必要です。出現回数を数えたりするだけなら、上記のコードで良いのでBinaryHeapを使う必要はありません。

必要なのは本当に優先度付きキューなのか？

結局はこういう問題な気がしてきました。追加した要素を昇順で格納する→優先度付きキュー→BinaryHeap、となってしまったので、データ構造をよく考える必要がありました。(と言っても、BinaryHeapの中身が見れないことを知らなかったのですが…)

というわけで、どんな時にどのデータ構造を使うのか検討しておいた方が良さそうです。

BinaryHeapを使うとき

値を追加したい。

常に最大(最小)の要素を出力したい。

最大(最小)の要素さえ出力できれば、他の要素は気にしない。

これが本当の「優先度付きキュー」ですね。

BTreeMapを使うとき

値を追加したい。

並び替えられた内部の状態を見たい。

keyの重複があっても、工夫して乗り切れる。