Rustで蟻本 〜データ構造〜

Rustで蟻本の学習を進めるシリーズ。データ構造の章に取り組みます。

プライオリティキュー

要素を追加(pushと呼ぶ)でき、最小の要素を取り出すこと(popと呼ぶ)ができるデータ構造をプライオリティキューと呼ぶ。抽象的な概念。例えばRustのVecは、要素の追加は容易である。さらに最小の要素を検索して取り出すことにすれば、上記のpush, popができ、プライオリティキューとして機能する。この場合、要素の追加(=push)にO(1)、最小の要素の取得(=pop)にO(n)の時間がかかる。

一方で要素を追加する際常に大きさ順にならべることにすると、要素の追加にはO(n)の時間がかかる(ランダムな配列に要素を挿入して並び替えるのにはO(nlogn)かかるが、整列された配列に一つ要素を挿入するのはO(n)で済む)が、最小の要素の取得はO(1)で済むようになる。まとめると、

• Vec1 = 最小要素取得時に検索 → 検索にはO(n)時間
• Vec2 = 要素追加時に大きさ順に並べる → 並び替えにO(n)時間

	push	pop
Vec1	O(1)	O(n)
Vec2	O(n)	O(1)

従って、要素を頻繁に追加する場合は、Vec1を使い、最小要素を頻繁に取得する場合はVec2を使えば良い。

ヒープ

ヒープはプライオリティキューを実現する一つのデータ構造であり、次の計算時間をもつ。

	push	pop
Heap	O(logn)	O(logn)

ちょうど(?)、Vec1, Vec2の中間の性質を持ったデータ構造である。詳しい解説は成書に譲るとして、Rustで実装してみたのを以下に示す。std::collections::BinaryHeapに倣ってpopは最小ではなく最大要素を取り出すようになっている。

collections.rs

#[derive(Default)]
pub struct Heap<T: PartialOrd> {
    v: Vec<T>,
}

impl<T: PartialOrd> Heap<T>
where
    Heap<T>: Default,
{
    pub fn new() -> Self {
        Self::default()
    }
    pub fn len(&self) -> usize {
        self.v.len()
    }
    //  最大要素を取り出す
    pub fn pop(&mut self) -> Option<T> {
        if self.len() > 0 {
            let top = self.v.swap_remove(0);
            let mut i = 0;
            while i * 2 + 1 < self.len() {
                let mut a = i * 2 + 1;
                if a + 1 < self.len() && self.v[a + 1] > self.v[a] {
                    a += 1; // aは子のうち大きい方のインデックス
                }
                if self.v[i] < self.v[a] {
                    self.v.swap(i, a);
                    i = a;
                } else {
                    break;
                }
            }
            Some(top)
        } else {
            None
        }
    }
    // 最大要素への参照を取得する
    pub fn peek(&self) -> &T {
        &self.v[0]
    }
    // 要素挿入
    pub fn push(&mut self, x: T) {
        self.v.push(x);
        let mut i = self.v.len() - 1;
        while i > 0 && self.v[i] > self.v[i >> 1] {
            self.v.swap(i, i >> 1);
            i >>= 1;
        }
    }
}

impl<T: PartialOrd> FromIterator<T> for Heap<T>
where
    Heap<T>: Default,
{
    fn from_iter<S>(iter: S) -> Self
    where
        S: IntoIterator<Item = T>,
    {
        let mut p = Self::default();
        for item in iter.into_iter() {
            p.push(item);
        }
        p
    }
}

Expedition

/// ガソリン補給問題
/// * `l` - 総距離
/// * `p` - 初期ガソリン
/// * `a` - ガソリンスタンドの位置
/// * `b` - ガソリンスタンドの補給可能量
fn expedition(l: i32, p: i32, mut a: Vec<i32>, mut b: Vec<i32>) -> i32 {
    a.push(l);
    b.push(0);
    let mut que = Heap::<i32>::new();
    let mut ans = 0;
    let mut pos = 0;
    let mut tank = p;
    for (&x, &g) in a.iter().zip(&b) {
        let d = x - pos; // 次に進む距離
        while tank < d {
            if let Some(q) = que.pop() {
                // O(logn)
                tank += q;
            } else {
                return -1;
            }
            ans += 1;
        }
        tank -= d;
        pos = x;
        que.push(g);
    }
    ans
}

Fence Repair

fn fence_repair(x: Vec<i32>) -> i32 {
    use std::cmp::Reverse;
    let mut ans = 0;
    let mut que: Heap<_> = x.into_iter().map(Reverse).collect();
    while que.len() >= 2 {
        // ループはn回
        let l = que.pop().unwrap().0 + que.pop().unwrap().0; // O(log n)
        ans += l;
        que.push(Reverse(l));
    }
    ans
}

メモ

• Heapで最小要素を取り出せるようにしたいときは、要素をstd::cmp::Reverseで包むと良い
• que.popがO(logn)なので、計算時間はO(nlogn)となった。

二分探索木

要素の追加、参照、削除ができるデータ構造である。
https://rust-unofficial.github.io/too-many-lists/
を参考にした。

struct Node<T: Ord> {
    elem: T,
    lhs: Link<T>,
    rhs: Link<T>,
}
struct Link<T: Ord>(Option<Box<Node<T>>>);

pub struct BinarySearchTree<T: Ord> {
    root: Link<T>,
}

impl<T: Ord> Link<T> {
    fn leaf(elem: T) -> Self {
        Link(Some(Box::new(Node {
            elem,
            lhs: Link(None),
            rhs: Link(None),
        })))
    }
    fn insert(&mut self, elem: T) {
        *self = match self.0.take() {
            None => Self::leaf(elem),
            Some(mut node) => {
                if node.elem < elem {
                    node.rhs.insert(elem);
                } else if elem < node.elem {
                    node.lhs.insert(elem);
                }
                Self(Some(node))
            }
        };
    }
    fn has(&self, elem: &T) -> bool {
        match self.0.as_deref() {
            None => false,
            Some(Node {
                elem: _elem,
                lhs,
                rhs,
            }) => _elem == elem || if elem < _elem { lhs } else { rhs }.has(elem),
        }
    }
    // 一番大きいやつをpop
    fn pop_first(&mut self) -> Option<T> {
        match self.0.take() {
            None => None,
            Some(mut b) => {
                if b.rhs.0.is_none() {
                    // 自分が最大
                    self.0 = b.lhs.0;
                    Some(b.elem) // elemを放出
                } else {
                    let ans = b.rhs.pop_first();
                    self.0 = Some(b);
                    ans
                }
            }
        }
    }
    // rootをpop
    fn pop_root(&mut self) -> Option<T> {
        match self.0.take() {
            None => None,
            Some(mut b) => {
                self.0 = if b.lhs.0.is_none() {
                    b.rhs.0
                } else {
                    Some(Box::new(Node {
                        elem: b.lhs.pop_first().unwrap(),
                        lhs: b.lhs,
                        rhs: b.rhs,
                    }))
                };
                Some(b.elem)
            }
        }
    }
    // 指定されたものをpop
    fn pop(&mut self, elem: &T) -> Option<T> {
        let mut link = self;
        loop {
            match &link.0.as_deref_mut() {
                None => {
                    return None;
                }
                Some(Node { elem: _elem, .. }) => {
                    if _elem == elem {
                        return link.pop_root();
                    }
                    link = if _elem < elem {
                        &mut link.0.as_mut().unwrap().rhs
                    } else {
                        &mut link.0.as_mut().unwrap().lhs
                    };
                }
            };
        }
    }
}

impl<T: Ord> BinarySearchTree<T> {
    pub fn new() -> Self {
        Self { root: Link(None) }
    }
    pub fn insert(&mut self, elem: T) {
        self.root.insert(elem);
    }
    pub fn has(&self, elem: &T) -> bool {
        self.root.has(elem)
    }
    pub fn pop(&mut self, elem: &T) -> Option<T> {
        self.root.pop(elem)
    }
}

#[cfg(test)]
mod test {
    use super::BinarySearchTree;
    #[test]
    fn insert_has_pop() {
        let mut bt = BinarySearchTree::new();
        bt.insert(1);
        bt.insert(3);
        bt.insert(2);
        bt.insert(4);
        assert!(bt.has(&4));
        assert!(bt.has(&3));
        assert!(bt.has(&2));
        assert!(bt.has(&1));

        assert!(!bt.has(&5));
        assert_eq!(bt.pop(&5), None);

        assert_eq!(bt.pop(&1), Some(1));
        assert_eq!(bt.pop(&3), Some(3));
        assert_eq!(bt.pop(&1), None);
        assert_eq!(bt.pop(&3), None);

        assert_eq!(bt.pop(&2), Some(2));
        assert_eq!(bt.pop(&2), None);

        bt.insert(1);
        bt.insert(1);
        assert_eq!(bt.pop(&1), Some(1));
        assert_eq!(bt.pop(&1), None);
    }
}

Union-Find木

いくつかの要素について、どのグループに属しているかがわかり、かつそれらの要素の属するグループ同士を統合できるデータ構造。木ではあるが、配列上に親の番号を記録していく形式で実装できるので、二分探索木などよりだいぶ簡単。

pub struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    pub fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
            rank: vec![0; n],
        }
    }
    pub fn find_parent(&mut self, x: usize) -> usize {
        if self.parent[x] == x {
            x
        } else {
            let par = self.find_parent(self.parent[x]);
            self.parent[x] = par;
            par
        }
    }
    pub fn unite(&mut self, mut x: usize, mut y: usize) {
        x = self.find_parent(x);
        y = self.find_parent(y);
        if x == y {
            return;
        }
        if self.rank[x] < self.rank[y] {
            self.parent[x] = y;
        } else {
            self.parent[y] = x;
            if self.rank[x] == self.rank[y] {
                self.rank[x] += 1;
            }
        }
    }
    pub fn same(&mut self, x: usize, y: usize) -> bool {
        self.find_parent(x) == self.find_parent(y)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn food_chain_test() {
        assert_eq!(
            3,
            food_chain(
                5,
                vec![
                    (1, 101, 1),
                    (2, 1, 2),
                    (2, 2, 3),
                    (2, 3, 3),
                    (1, 1, 3),
                    (2, 3, 1),
                    (1, 5, 5),
                ]
            )
        );
    }
    fn food_chain(n: usize, informations: Vec<(usize, usize, usize)>) -> usize {
        let mut uf = UnionFind::new(n * 3);
        let mut ans = 0;
        for (t, mut x, mut y) in informations {
            // 1<=x<=n, 1<=y<=nを確認
            if x <= 0 || n < x || y <= 0 || n < y {
                ans += 1;
                continue;
            }
            // 0, ..., n-1に合わせる
            x -= 1;
            y -= 1;
            if t == 1 {
                // 矛盾を見つけるときは、対称性よりxはグループAに属していると仮定して良い
                // その場合、yがBまたはCに属するのが矛盾である
                if uf.same(x, y + n) || uf.same(x, y + 2 * n) {
                    ans += 1;
                } else {
                    uf.unite(x, y);
                    uf.unite(x + n, y + n);
                    uf.unite(x + 2 * n, y + 2 * n);
                }
            } else {
                // 同様にxはグループAに属していると仮定
                // yがAまたはCに属するのが矛盾
                if uf.same(x, y) || uf.same(x, y + 2 * n) {
                    ans += 1;
                } else {
                    uf.unite(x, y + n);
                    uf.unite(x + n, y + 2 * n);
                    uf.unite(x + 2 * n, y);
                }
            }
        }
        ans
    }
}