multiset/
lib.rs

1//! A multiset that allows multiple occurrences of the same element.
2//!
3//! Backed by a [`std::collections::BTreeMap`], so elements are kept in sorted
4//! order by key. `BTreeMap` is preferred over `HashMap` here to avoid
5//! collision blowups and to allow `O(log n)` `first` / `last` queries.
6//!
7//! All operations are `O(log n)`; `len` and `is_empty` are `O(n)` (they sum
8//! all counts) but can be cached if needed.
9
10use std::collections::BTreeMap;
11
12/// A multiset of values of type `T` (any `Ord + Clone`).
13#[derive(Debug, Default)]
14pub struct MultiSet<T> {
15    elems: BTreeMap<T, usize>,
16}
17
18impl<T: Ord + Clone> MultiSet<T> {
19    /// Construct an empty multiset.
20    pub fn new() -> Self {
21        MultiSet {
22            elems: BTreeMap::new(),
23        }
24    }
25
26    /// Insert one more occurrence of `el`.
27    pub fn insert(&mut self, el: T) {
28        *self.elems.entry(el).or_insert(0) += 1;
29    }
30
31    /// Remove one occurrence of `el`. Returns `false` if `el` was not present.
32    pub fn remove(&mut self, el: &T) -> bool {
33        if let Some(count) = self.elems.get_mut(el) {
34            if *count > 1 {
35                *count -= 1;
36            } else {
37                self.elems.remove(el);
38            }
39            true
40        } else {
41            false
42        }
43    }
44
45    /// Remove all occurrences of `el`. Returns `false` if `el` was not present.
46    pub fn remove_all(&mut self, el: &T) -> bool {
47        self.elems.remove(el).is_some()
48    }
49
50    /// Whether `el` is present at least once.
51    pub fn contains(&self, el: &T) -> bool {
52        self.elems.contains_key(el)
53    }
54
55    /// Smallest element (by `Ord`), or `None` if the multiset is empty.
56    pub fn first(&self) -> Option<T> {
57        self.elems.keys().next().cloned()
58    }
59
60    /// Largest element (by `Ord`), or `None` if the multiset is empty.
61    pub fn last(&self) -> Option<T> {
62        self.elems.keys().next_back().cloned()
63    }
64
65    /// Iterate over `(element, count)` pairs in ascending order.
66    pub fn iter(&self) -> impl Iterator<Item = (&T, &usize)> {
67        self.elems.iter()
68    }
69
70    /// Consume the multiset and iterate over `(element, count)` pairs in
71    /// ascending order.
72    pub fn into_iter(self) -> impl Iterator<Item = (T, usize)> {
73        self.elems.into_iter()
74    }
75
76    /// Iterate mutably over `(element, count)` pairs in ascending order.
77    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&T, &mut usize)> {
78        self.elems.iter_mut()
79    }
80
81    /// Number of occurrences of `el` (0 if absent).
82    pub fn count(&self, el: &T) -> usize {
83        *self.elems.get(el).unwrap_or(&0)
84    }
85
86    /// Total number of elements (with multiplicities).
87    pub fn len(&self) -> usize {
88        self.elems.values().sum()
89    }
90
91    /// Whether the multiset has no elements.
92    pub fn is_empty(&self) -> bool {
93        self.elems.is_empty()
94    }
95
96    /// Remove all elements.
97    pub fn clear(&mut self) {
98        self.elems.clear();
99    }
100}
101
102impl<T: Ord + Clone> FromIterator<T> for MultiSet<T> {
103    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
104        let mut multiset = MultiSet::new();
105        for el in iter {
106            multiset.insert(el);
107        }
108        multiset
109    }
110}
multiset/lib.rs

multiset/
lib.rs
