nekolib/ds/
btree_multiset.rs

1//! 多重集合。
2
3use std::collections::{btree_map::Iter as BTreeMapIter, BTreeMap};
4use std::fmt::{self, Debug};
5
6/// 多重集合。
7#[derive(Clone, Eq, PartialEq)]
8pub struct BTreeMultiset<K>(BTreeMap<K, usize>, usize);
9
10impl<K: Ord> BTreeMultiset<K> {
11    pub fn new() -> Self { Self(BTreeMap::new(), 0) }
12    pub fn insert(&mut self, k: K) { self.insert_n(k, 1) }
13    pub fn insert_n(&mut self, k: K, n: usize) {
14        if n == 0 {
15            return;
16        }
17        *self.0.entry(k).or_insert(0) += n;
18        self.1 += n;
19    }
20    pub fn remove(&mut self, k: &K) { self.remove_n(k, 1) }
21    pub fn remove_n(&mut self, k: &K, n: usize) {
22        if n == 0 {
23            return;
24        }
25        match self.0.get(k) {
26            Some(&c) if c <= n => {
27                self.0.remove(&k);
28                self.1 -= c;
29            }
30            Some(_) => {
31                *self.0.get_mut(k).unwrap() -= n;
32                self.1 -= n;
33            }
34            None => {}
35        }
36    }
37    pub fn min(&self) -> Option<&K> { self.0.keys().next() }
38    pub fn max(&self) -> Option<&K> { self.0.keys().next_back() }
39    pub fn is_empty(&self) -> bool { self.0.is_empty() }
40    pub fn len(&self) -> usize { self.1 }
41    pub fn count(&self, k: &K) -> usize { *self.0.get(k).unwrap_or(&0) }
42
43    pub fn iter(&self) -> Iter<'_, K> { Iter::new(self) }
44}
45
46pub struct Iter<'a, K: 'a> {
47    iter: BTreeMapIter<'a, K, usize>,
48    key: Option<&'a K>,
49    count: usize,
50}
51
52impl<'a, K: 'a + Ord> Iter<'a, K> {
53    pub fn new(ms: &'a BTreeMultiset<K>) -> Self {
54        let mut iter = ms.0.iter();
55        if let Some((key, &count)) = iter.next() {
56            Self { iter, key: Some(key), count }
57        } else {
58            Self { iter, key: None, count: 0 }
59        }
60    }
61}
62
63impl<'a, K: 'a> Iterator for Iter<'a, K> {
64    type Item = &'a K;
65    fn next(&mut self) -> Option<&'a K> {
66        if self.key.is_none() {
67            return None;
68        }
69
70        if self.count > 0 {
71            self.count -= 1;
72        } else if let Some((k, &v)) = self.iter.next() {
73            self.key = Some(k);
74            self.count = v - 1;
75        } else {
76            self.key = None;
77        }
78        self.key
79    }
80}
81
82impl<'a, K: Ord> IntoIterator for &'a BTreeMultiset<K> {
83    type Item = &'a K;
84    type IntoIter = Iter<'a, K>;
85    fn into_iter(self) -> Self::IntoIter { self.iter() }
86}
87
88impl<K: Debug> Debug for BTreeMultiset<K> {
89    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
90        f.debug_map().entries(self.0.iter()).finish()
91    }
92}
93
94#[test]
95fn test_debug_fmt() {
96    let mut ms = BTreeMultiset::new();
97    ms.insert(5);
98    ms.insert(5);
99    ms.insert(7);
100    ms.insert_n(4, 0);
101    ms.remove_n(&3, 0);
102    assert_eq!(format!("{:?}", ms), "{5: 2, 7: 1}");
103}
104
105#[test]
106fn test_iter() {
107    let mut ms = BTreeMultiset::new();
108    ms.insert(5);
109    ms.insert(5);
110    ms.insert(6);
111    ms.insert(6);
112    ms.insert(6);
113    ms.insert(7);
114    ms.insert_n(4, 0);
115    ms.remove_n(&3, 0);
116    assert_eq!(ms.iter().copied().collect::<Vec<_>>(), [5, 5, 6, 6, 6, 7]);
117}
118
119#[test]
120fn test_count() {
121    let mut ms = BTreeMultiset::new();
122    assert_eq!(ms.len(), 0);
123
124    ms.insert_n(5, 2);
125    assert_eq!(ms.len(), 2);
126    assert_eq!(ms.count(&5), 2);
127    assert_eq!(ms.count(&4), 0);
128
129    ms.remove_n(&5, 3);
130    assert_eq!(ms.len(), 0);
131    assert_eq!(ms.count(&5), 0);
132    assert_eq!(ms.count(&4), 0);
133}
nekolib/ds/btree_multiset.rs

nekolib/ds/
btree_multiset.rs
