【Rust】優先度付きキュー

優先度付きキューとは

優先度付きキュー：優先度の高いものから順に取り出すキュー

RustではBinaryHeapを利用して実装する。BinaryHeapはヒープと呼ばれるデータ構造を持つので以下で簡単に触れておく。

ヒープ

ヒープとはここでは木の一種であり、以下を満たす：

• 親子ノード間に大小関係が定義されている
• 完全二分木（最下層が左から順に、それ以外は完全に埋まっている二分木）

最大ヒープ

• 親ノード >= 子ノード
• 根が最大の要素

最小ヒープ

• 親ノード <= 子ノード
• 根が最小の要素

要素が追加されると、その要素はまず最下層に配置され、親子の大小関係が正しくなるまで親ノードと入れ替えを繰り返す。
ヒープの根が最も優先度の高い要素に当たる。

Rustで優先度付きキューを実装する

先ほども触れたようにRustではBinaryHeapを利用して優先度付きキューを実装できる。

BinaryHeap in std::collections - Rust

BinaryHeapは最大ヒープであることに注意。

主なメソッド

• push()：要素を追加
• peek()：優先度が最も高い要素を参照
• pop()：優先度が最も高い要素を取り出し（削除+取得）
• len()：要素数を取得
• is_empty()：空かどうかを取得
• clear()：空にする

最大ヒープ

降順に並べたい場合は最大ヒープ。

let mut max_heap: BinaryHeap<i32> = BinaryHeap::new();
max_heap.push(1);
max_heap.push(2);
max_heap.push(3);
if let Some(x) = max_heap.peek() {
    println!("{}", x); // 3
}
if let Some(x) = max_heap.pop() {
    println!("{}", x);  // 3
}

最小ヒープ

昇順に並べたい場合は最小ヒープ。
最小ヒープを実現するには core::cmp::Reverseを用いる。

let mut min_heap: BinaryHeap<Reverse<i32>> = BinaryHeap::new();
min_heap.push(Reverse(1));
min_heap.push(Reverse(2));
min_heap.push(Reverse(3));
if let Some(Reverse(x)) = min_heap.peek() {
    println!("{}", x);  // 1
}
if let Some(Reverse(x)) = min_heap.pop() {
    println!("{}", x);  // 1
}

タプルを格納する

タプルは1番目の要素から順に比較され、大小関係が決定されるらしい。

let mut heap: BinaryHeap<Reverse<(i32, &str)>> = BinaryHeap::new();
heap.push(Reverse((2, "aaa")));  // 1
heap.push(Reverse((2, "bbb")));  // 2
heap.push(Reverse((1, "xxx")));  // 3
while let Some(Reverse(x)) = heap.pop() {
    println!("{} {}", x.0, x.1)  // 3, 1, 2 の順で出力される
}

カスタム優先度

Ordトレイトを実装することで優先度を決定するロジックを自前で用意できる。

use std::{
    cmp::{Ordering, Reverse},
    collections::BinaryHeap,
};

fn main() {
    let mut heap: BinaryHeap<Reverse<User>> = BinaryHeap::new();
    heap.push(Reverse(User::new("Alice", 25))); // 1
    heap.push(Reverse(User::new("Bob", 20))); // 2
    heap.push(Reverse(User::new("Cassie", 20))); // 3
    while let Some(Reverse(user)) = heap.pop() {
        println!("{:?}", user); // 3, 2, 1 の順番で出力される
    }
}

#[derive(Debug, PartialEq, Eq)]
struct User {
    name: String,
    age: i32,
}

impl User {
    fn new(name: &str, age: i32) -> Self {
        Self {
            name: name.to_string(),
            age: age,
        }
    }
}

impl PartialOrd for User {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(&other))
    } 
}

impl Ord for User {
    fn cmp(&self, other: &Self) -> Ordering {
        // age: 昇順、name: 長さの降順 の順番
        self.age.cmp(&other.age).then(other.name.len().cmp(&self.name.len()))
    }
}

参考文献

ヒープとは - 意味をわかりやすく - IT用語辞典 e-Words
Rustで競技プログラミングよくばりセット #AtCoder - Qiita