RustでMultiSetを実装する

TL;DR

• BTreeMapをベースにRustでマルチセットを作りました
• C++にあるMultiSetとだいたい同じことができます
• イテレータは実装していません (追記：実装しました)

使った報告をしてくれると僕が泣いて喜びます。

モチベーション

AtCoderを最近は頑張ってやってるんですが、ABC308F: Vouchersなどの問題でMultiSetが欲しいな〜みたいなタイミングがちょこちょこあります。”MultiSetがないので、実装きついです”は良くないので、コンテスト中に使えるように処理をまとめよう！ということで書きました。

MultiSetとは？

MultiSetは数学における多重集合に対応するデータ構造です。普通、集合といえば要素の重複が許されませんが、多重集合は重複を許します。イメージとしては、順序のない数列のような感じです。

競技プログラミングでもよく使われるデータ構造ですが、C++くらいでしか標準で実装されているのを寡聞にして聞いたことがないので、Python勢やRust勢のみなさんどうやって通してるんでしょうか...？

実装

BTreeMapに重複をカウントして実装しています。値を挿入すると、BTreeMap上のカウントを1増やす、BTreeMap上のカウントが0になったらkeyを消すという単純な方法です。

普通MultiSetといえば、赤黒木などの平衝二分木を用いて実装されます。BTreeMapでの実装は重複が大量にある場合有利ですが、多くの場合においては平衝二分木のほうが良いと思います（未検証）。

イテレータの実装はいい感じに私の技量ではできなかったので、今回は実装していません。単に、重複を削除したものを返しても良かったんですが、イテレータじゃないよなーと思ったのでそのままです。

いい感じに実装できたのでイテレータも実装しました。

BTreeMapを使うのでOrdトレイトが実装されている必要があります。

コードは長いので折りたたみ

コード

#[allow(dead_code)]
mod procon_lib {
    use std::{
        borrow::Borrow,
        collections::{
            btree_map::{self},
            BTreeMap,
        },
        ops::{Bound, RangeBounds},
    };

    #[derive(Debug, PartialEq, Clone)]
    pub struct MultiSet<T> {
        size: usize,
        btree_map: BTreeMap<T, usize>,
    }

    impl<T: Ord> Default for MultiSet<T> {
        fn default() -> Self {
            Self::new()
        }
    }

    impl<T: Ord> From<Vec<T>> for MultiSet<T> {
        fn from(value: Vec<T>) -> Self {
            let size = value.len();

            let mut btree_map = BTreeMap::default();
            for key in value {
                *btree_map.entry(key).or_insert(0) += 1;
            }

            Self { size, btree_map }
        }
    }

    impl<T: Ord> MultiSet<T> {
        pub fn clear(&mut self) {
            self.size = 0;
            self.btree_map.clear();
        }

        pub fn contains<Q>(&self, value: &Q) -> bool
        where
            T: Borrow<Q>,
            Q: Ord + ?Sized,
        {
            self.btree_map.contains_key(value)
        }

        pub fn insert(&mut self, value: T) {
            self.size += 1;
            *self.btree_map.entry(value).or_insert(0) += 1;
        }

        pub fn first(&self) -> Option<&T> {
            if let Some((key, _)) = self.btree_map.iter().next() {
                Some(key)
            } else {
                None
            }
        }

        pub fn last(&self) -> Option<&T> {
            if let Some((key, _)) = self.btree_map.iter().next_back() {
                Some(key)
            } else {
                None
            }
        }

        pub fn is_empty(&self) -> bool {
            self.btree_map.is_empty()
        }

        pub fn len(&self) -> usize {
            self.btree_map.len()
        }

        pub fn size(&self) -> usize {
            self.size
        }

        pub fn marge(&mut self, other: &mut MultiSet<T>)
        where
            T: Clone,
        {
            self.size += other.size;

            for (key, val) in other.btree_map.iter() {
                if let Some(prev) = self.btree_map.get_mut(key) {
                    *prev += *val;
                } else {
                    self.btree_map.insert(key.clone(), *val);
                }
            }
        }

        pub fn new() -> Self {
            Self {
                size: 0,
                btree_map: BTreeMap::new(),
            }
        }

        pub fn pop_first(&mut self) -> Option<T>
        where
            T: Clone,
        {
            if self.is_empty() {
                None
            } else {
                self.size -= 1;

                let first = self.first().unwrap().clone();
                self.remove(&first);
                Some(first)
            }
        }

        pub fn pop_last(&mut self) -> Option<T>
        where
            T: Clone,
        {
            if self.is_empty() {
                None
            } else {
                self.size -= 1;

                let last = self.last().unwrap().clone();
                self.remove(&last);
                Some(last)
            }
        }

        pub fn remove(&mut self, value: &T) -> bool
        where
            T: Clone,
        {
            self.btree_map.entry(value.clone()).and_modify(|e| *e -= 1);
            if let Some(&cnt) = self.btree_map.get(&value) {
                if cnt == 0 {
                    self.btree_map.remove(&value);
                }

                self.size -= 1;
                true
            } else {
                false
            }
        }

        pub fn lower_bound<Q>(&self, bound: Bound<&Q>) -> Option<&T>
        where
            T: Borrow<Q>,
            Q: Ord,
        {
            match bound {
                Bound::Unbounded => unreachable!(),
                _ => {
                    if let Some((key, _)) = self.btree_map.range((bound, Bound::Unbounded)).next() {
                        Some(key)
                    } else {
                        None
                    }
                }
            }
        }

        pub fn upper_bound<Q>(&self, bound: Bound<&Q>) -> Option<&T>
        where
            T: Borrow<Q>,
            Q: Ord,
        {
            match bound {
                Bound::Unbounded => unreachable!(),
                _ => {
                    if let Some((key, _)) =
                        self.btree_map.range((Bound::Unbounded, bound)).next_back()
                    {
                        Some(key)
                    } else {
                        None
                    }
                }
            }
        }

        pub fn iter(&self) -> Iter<'_, T> {
            Iter {
                range: self.range(..),
            }
        }

        pub fn range<U: ?Sized, R>(&self, range: R) -> Range<'_, T>
        where
            U: Ord,
            T: Borrow<U> + Ord,
            R: RangeBounds<U>,
        {
            Range {
                last: None,
                counter: 0,
                range: self.btree_map.range(range),
            }
        }

        pub fn count<Q>(&self, value: &Q) -> usize
        where
            T: Borrow<Q>,
            Q: Ord,
        {
            if let Some(&cnt) = self.btree_map.get(value) {
                cnt
            } else {
                0
            }
        }
    }

    #[derive(Debug, Clone, Default)]
    pub struct Range<'a, T>
    where
        T: 'a,
    {
        last: Option<&'a T>,
        counter: usize,
        range: btree_map::Range<'a, T, usize>,
    }

    impl<'a, T> Iterator for Range<'a, T> {
        type Item = &'a T;
        fn next(&mut self) -> Option<Self::Item> {
            if self.counter == 0 {
                if let Some((elem, &cnt)) = self.range.next() {
                    self.last = Some(elem);
                    self.counter = cnt - 1;
                    Some(elem)
                } else {
                    None
                }
            } else {
                self.counter -= 1;
                self.last
            }
        }
    }

    impl<'a, T> DoubleEndedIterator for Range<'a, T> {
        fn next_back(&mut self) -> Option<Self::Item> {
            if self.counter == 0 {
                if let Some((elem, &cnt)) = self.range.next_back() {
                    self.last = Some(elem);
                    self.counter = cnt;
                    Some(elem)
                } else {
                    None
                }
            } else {
                self.counter -= 1;
                self.last
            }
        }
    }

    #[derive(Clone, Debug, Default)]
    pub struct Iter<'a, T>
    where
        T: 'a,
    {
        range: Range<'a, T>,
    }

    impl<'a, T> Iterator for Iter<'a, T> {
        type Item = &'a T;
        fn next(&mut self) -> Option<Self::Item> {
            self.range.next()
        }
    }

    impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
        fn next_back(&mut self) -> Option<Self::Item> {
            self.range.next_back()
        }
    }
}

機能と使い方

new(), default()

インスタンスを生成します。どちらも等価です。

use procon_lib::MultiSet;

fn main() {
    // どちらでもいっしょ
    let multiset = MultiSet::new();
    let multiset = MultiSet::default();
}

from(Vec<T>)

Vec<T>からインスタンスを生成します。

use procon_lib::MultiSet;

fn main() {
    let multiset = MultiSet::from(vec![1, 1, 2, 3]);
}

clear()

MultiSetを空集合にします。

use procon_lib::MultiSet;

fn main() {
    let mut a = MultiSet::from(vec![1, 2, 3]);

    a.clear();
    assert_eq!(a, MultiSet::from(vec![]));
}

contains(&T)

MultiSetが引数を元として含むかどうかを判定します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 2, 2, 3]);

    assert_eq!(a.contains(&1), true);
    assert_eq!(a.contains(&0), false);
}

insert(T)

MultiSetに引数を追加します。

use procon_lib::MultiSet;

fn main() {
    let mut a = MultiSet::from(vec![1, 2, 2, 3]);

    a.insert(0);
    assert_eq!(a, MultiSet::from(vec![0, 1, 2, 2, 3]));
}

first(), last()

firstは最小の要素を、last()は最大の要素を取得します。

use procon_lib::MultiSet;

fn main() {
    let mut a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.first(), Some(&1));
    assert_eq!(a.last(), Some(&3));
}

pop_first, pop_last()

pop_first()は最小の要素を、pop_last()は最大の要素を1つ削除します。

use procon_lib::MultiSet;

fn main() {
    let mut a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.pop_first(), Some(1));
    assert_eq!(a, MultiSet::from(vec![1, 2, 2, 3]));

    assert_eq!(a.pop_last(), Some(3));
    assert_eq!(a, MultiSet::from(vec![1, 2, 2]));
}

is_empty(), len(), size()

is_empty()は空集合であるかどうかを判定します。len()は集合の元の種類を返す一方、size()は元の個数を返します（命名が怪しい）。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.is_empty(), false);
    assert_eq!(a.len(), 3);
    assert_eq!(a.size(), 5);
}

count(&Q)

引数の要素が集合内にいくつあるのかを返します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.count(&2), 2);
}

remove(&T)

MultiSetから引数の元を一つ削除します。削除に成功すればtrueを、削除できなければfalseを返します。

fn main() {
    let mut a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.remove(&2), true);
    assert_eq!(a.remove(&4), false);

    assert_eq!(a, MultiSet::from(vec![1, 1, 2, 3]));
}

lower_bound(Bound<&Q>)

std::ops::Boundを引数にとり、引数よりも大きい、または引数以上のものの中から最小のものを取り出します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.lower_bound(std::ops::Bound::Included(&0)), Some(&1));
    assert_eq!(a.lower_bound(std::ops::Bound::Included(&2)), Some(&2));
    assert_eq!(a.lower_bound(std::ops::Bound::Excluded(&2)), Some(&3));
    assert_eq!(a.lower_bound(std::ops::Bound::Excluded(&4)), None);
}

upper_bound(Bound<&Q>)

std::ops::Boundを引数にとり、引数よりも小さい、または引数以下のものの中から最大のものを取り出します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    assert_eq!(a.upper_bound(std::ops::Bound::Excluded(&0)), None);
    assert_eq!(a.upper_bound(std::ops::Bound::Excluded(&2)), Some(&1));
    assert_eq!(a.upper_bound(std::ops::Bound::Included(&2)), Some(&2));
    assert_eq!(a.upper_bound(std::ops::Bound::Included(&4)), Some(&3));
}

iter()

参照のイテレータを返します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 1, 2, 2, 3]);

    let mut b = vec![];
    for &e in a.iter() {
        b.push(e);
    }

    assert_eq!(b, vec![1, 1, 2, 2, 3]);
    assert_eq!(a, MultiSet::from(b));
}

range(..)

引数の範囲にある要素の参照のIterableを返します。

use procon_lib::MultiSet;

fn main() {
    let a = MultiSet::from(vec![1, 2, 3, 4, 4, 5, 6, 7, 8]);

    let mut b = vec![];
    for &e in a.range(2..=5) {
        b.push(e);
    }

    assert_eq!(b, vec![2, 3, 4, 4, 5]);
}

marge()

片方のMultiSetをもう片方にマージして、和集合を生成します。

use procon_lib::MultiSet;

fn main() {
    let mut a = MultiSet::from(vec![1, 2, 3]);
    let mut b = MultiSet::from(vec![1, 1, 2, 3]);
    
    a.marge(&mut b);
    assert_eq!(a, MultiSet::from(vec![1, 1, 1, 2, 2, 3, 3]));
}

使用例

• AtCoder Beginner Contest 308: F - Vouchers
• AtCoder Beginner Contest 253: C - Max - Min Query
• AtCoder Beginner Contest 170: E - Smart Infants