読者です 読者をやめる 読者になる 読者になる

Rustで関数型っぽい見た目のペアリングヒープを実装

Rust

Rustでペアリングヒープを実装してみたのでコードをメモとして残しておく。
無駄に関数型っぽい見た目の実装(selfを使わずに更新結果を戻り値で返す)となっている。
最低限のテストは書いたが正しく動作するかどうかは自信がない。
効率周りも無頓着。

// $ rustc --version
// rustc 1.0.0 (a59de37e9 2015-05-13) (built 2015-05-14)
//
// $ cargo new pairing_heaps
// $ cd pairing_heaps
// $ emacs src/lib.rs
// $ cargo build

// ヒープ型
pub enum PairingHeap<T> {
    Empty,
    Node(T, Vec<PairingHeap<T>>),
}

// 空ヒープ生成
pub fn new<T>() -> PairingHeap<T> {
    PairingHeap::Empty
}

// 単一要素ヒープ生成
pub fn singleton<T>(item: T) -> PairingHeap<T> {
    PairingHeap::Node(item, Vec::new())
}

// 空判定
pub fn is_empty<T>(h: &PairingHeap<T>) -> bool {
    match h {
        &PairingHeap::Empty => true,
        _                   => false,
    }
}

// 要素数取得
pub fn size<T>(h: &PairingHeap<T>) -> usize {
    match h {
        &PairingHeap::Empty           => 0,
        &PairingHeap::Node(_, ref hs) => hs.iter().fold(1, |sum, child| sum + size(child))
    }
}

// 要素追加
pub fn push<T: PartialOrd>(item: T, h: PairingHeap<T>) -> PairingHeap<T> {
    merge(singleton(item), h)
}

// 最小要素取り出し
pub fn pop<T: PartialOrd>(h: PairingHeap<T>) -> (Option<T>, PairingHeap<T>) {
    match h {
        PairingHeap::Empty       => (None, new()),
        PairingHeap::Node(x, xs) => (Some(x), merge_pairs(xs)),
    }
}

// ヒープ同士のマージ
pub fn merge<T: PartialOrd>(h1: PairingHeap<T>, h2: PairingHeap<T>) -> PairingHeap<T> {
    match (h1, h2) {
        (PairingHeap::Empty, h2) => h2,
        (h1, PairingHeap::Empty) => h1,
        (PairingHeap::Node(x, mut xs), PairingHeap::Node(y, mut ys)) =>
            if x < y {
                xs.push(PairingHeap::Node(y, ys));
                PairingHeap::Node(x, xs)
            } else {
                ys.push(PairingHeap::Node(x, xs));
                PairingHeap::Node(y, ys)
            }
    }
}

fn merge_pairs<T: PartialOrd>(mut hs: Vec<PairingHeap<T>>) -> PairingHeap<T> {
    match hs.len() {
        0 => PairingHeap::Empty,
        1 => hs.pop().unwrap(),
        _ => {
            let h1 = hs.pop().unwrap();
            let h2 = hs.pop().unwrap();
            merge(merge(h1, h2), merge_pairs(hs))
        }
    }
}

残りはテスト関数:

// $ emacs src/lib.rs
// $ cargo test
#[test]
fn it_works() {
    // 空ヒープ生成
    let h: PairingHeap<u8> = new();
    assert!(is_empty(&h));
    assert_eq!(0, size(&h));

    // 要素追加1
    let h = push(2, h);
    assert_eq!(1, size(&h));

    // 要素追加2
    let h = push(1, h);
    assert_eq!(2, size(&h));

    // 最小値取り出し1
    let (x, h) = pop(h);
    assert_eq!(Some(1), x);
    assert_eq!(1, size(&h));

    // 最小値取り出し2
    let (x, h) = pop(h);
    assert_eq!(Some(2), x);
    assert_eq!(0, size(&h));

    // 空ヒープからの最小値取り出し
    let (x, h) = pop(h);
    assert_eq!(None, x);
    assert_eq!(0, size(&h));
}