A set structure which contains multiple same items with count of them, implemented for competitive programming in Rust.

btree_multi_set.rs

use std::{borrow::Borrow, collections::BTreeMap, num::NonZeroUsize};

#[derive(Debug, Clone)]
pub struct BTreeMultiSet<T> {
    counts: BTreeMap<T, NonZeroUsize>,
    len: usize,
}

impl<T> BTreeMultiSet<T> {
    pub const fn new() -> Self {
        Self {
            counts: BTreeMap::new(),
            len: 0,
        }
    }

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

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

    pub fn count<Q>(&self, item: &Q) -> usize
    where
        Q: ?Sized + Ord,
        T: Borrow<Q> + Ord,
    {
        self.counts.get(item).map_or(0, |count| count.get())
    }

    pub fn insert(&mut self, item: T)
    where
        T: Ord,
    {
        self.counts
            .entry(item)
            .and_modify(|count| *count = count.checked_add(1).unwrap())
            .or_insert_with(|| NonZeroUsize::new(1).unwrap());
        self.len += 1;
    }

    pub fn remove<Q>(&mut self, item: &Q) -> bool
    where
        Q: ?Sized + Ord,
        T: Borrow<Q> + Ord,
    {
        if self.counts.get(item) == Some(&NonZeroUsize::new(1).unwrap()) {
            self.counts.remove(item);
            self.len -= 1;
            true
        } else if let Some(count) = self.counts.get_mut(item) {
            *count = NonZeroUsize::new(count.get() - 1).unwrap();
            self.len -= 1;
            true
        } else {
            false
        }
    }

    pub fn min(&self) -> Option<&T>
    where
        T: Ord,
    {
        self.counts.first_key_value().map(|(item, _)| item)
    }

    pub fn max(&self) -> Option<&T>
    where
        T: Ord,
    {
        self.counts.last_key_value().map(|(item, _)| item)
    }
}