• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 #![cfg(feature = "std")]
2 use core::hash::Hash;
3 use std::collections::HashMap;
4 
5 /// Multiset container inspired by Python's `collections.Counter`.
6 pub struct Counter<K> {
7     map: HashMap<K, usize>,
8 }
9 
10 impl<K> Counter<K>
11 where
12     K: Hash + Eq,
13 {
14     /// make an empty Counter
new() -> Counter<K>15     pub fn new() -> Counter<K> {
16         Counter {
17             map: HashMap::new(),
18         }
19     }
20 
21     /// Create a counter from a sequence.
from_iter<I>(iter: I) -> Counter<K> where I: IntoIterator<Item = K>,22     pub fn from_iter<I>(iter: I) -> Counter<K>
23     where
24         I: IntoIterator<Item = K>,
25     {
26         let mut counter = Counter::new();
27         counter.update(iter);
28         counter
29     }
30 
31     /// Merge items from a sequence into the Counter
update<I>(&mut self, iter: I) where I: IntoIterator<Item = K>,32     pub fn update<I>(&mut self, iter: I)
33     where
34         I: IntoIterator<Item = K>,
35     {
36         for item in iter {
37             let entry = self.map.entry(item).or_insert(0);
38             *entry += 1;
39         }
40     }
41 
42     // How many items there are in total in the Counter.
count(&self) -> usize43     pub fn count(&self) -> usize {
44         self.map.values().sum()
45     }
46 
47     /// Unique elements in the set
keys(&self) -> impl Iterator<Item = &K>48     pub fn keys(&self) -> impl Iterator<Item = &K> {
49         self.map.keys()
50     }
51 
52     /// The count of things without the things
values(&self) -> impl Iterator<Item = &usize>53     pub fn values(&self) -> impl Iterator<Item = &usize> {
54         self.map.values()
55     }
56 
57     /// .
get(&self, key: &K) -> Option<&usize>58     pub fn get(&self, key: &K) -> Option<&usize> {
59         self.map.get(key)
60     }
61 
62     /// Merge two counters together.
merge<'a>(&'a self, rhs: &'a Counter<K>) -> Counter<&'a K>63     pub fn merge<'a>(&'a self, rhs: &'a Counter<K>) -> Counter<&'a K> {
64         let mut result: HashMap<&K, usize> = HashMap::new();
65         for (key, lhs_count) in &self.map {
66             let rhs_count = rhs.map.get(key).unwrap_or(&0);
67             result.insert(key, *lhs_count + rhs_count);
68         }
69         for (key, rhs_count) in &rhs.map {
70             if !self.map.contains_key(key) {
71                 result.insert(key, *rhs_count);
72             }
73         }
74         Counter { map: result }
75     }
76 
77     /// How many there are common items in the given multisets.
intersect_count(&self, rhs: &Counter<K>) -> usize78     pub fn intersect_count(&self, rhs: &Counter<K>) -> usize {
79         let mut result = 0;
80         for (key, lhs_count) in &self.map {
81             if let Some(rhs_count) = rhs.map.get(key) {
82                 result += lhs_count.min(rhs_count);
83             }
84         }
85         result
86     }
87 
88     /// How many there are items in total in both multisets.
union_count(&self, rhs: &Counter<K>) -> usize89     pub fn union_count(&self, rhs: &Counter<K>) -> usize {
90         let mut result = 0;
91         for (key, lhs_count) in &self.map {
92             let rhs_count = rhs.map.get(key).unwrap_or(&0);
93             result += lhs_count.max(rhs_count);
94         }
95         for (key, rhs_count) in &rhs.map {
96             if !self.map.contains_key(key) {
97                 result += rhs_count;
98             }
99         }
100         result
101     }
102 
103     /// How many there are item in left that aren't in the right
diff_count(&self, rhs: &Counter<K>) -> usize104     pub fn diff_count(&self, rhs: &Counter<K>) -> usize {
105         let mut result = 0;
106         for (key, lhs_count) in &self.map {
107             let rhs_count = rhs.map.get(key).unwrap_or(&0);
108             if lhs_count > rhs_count {
109                 result += lhs_count - rhs_count;
110             }
111         }
112         result
113     }
114 }
115 
116 #[cfg(test)]
117 mod tests {
118     use super::*;
119     use assert2::assert;
120     use rstest::rstest;
121 
eq<K: Hash + Eq>(lhs: &Counter<K>, rhs: &Counter<K>) -> bool122     pub fn eq<K: Hash + Eq>(lhs: &Counter<K>, rhs: &Counter<K>) -> bool {
123         for (key, lhs_count) in &lhs.map {
124             if let Some(rhs_count) = rhs.map.get(key) {
125                 if lhs_count != rhs_count {
126                     return false;
127                 }
128             } else {
129                 return false;
130             }
131         }
132         for (key, rhs_count) in &rhs.map {
133             if let Some(lhs_count) = lhs.map.get(key) {
134                 if lhs_count != rhs_count {
135                     return false;
136                 }
137             } else {
138                 return false;
139             }
140         }
141         true
142     }
143 
144     #[rstest]
smoke()145     fn smoke() {
146         let c1 = Counter::from_iter(1..=5);
147         let c2 = Counter::from_iter(3..=7);
148         assert!(eq(&c1, &c1));
149         assert!(!eq(&c1, &c2));
150         // assert!(eq(c1.intersect(&c2), &Counter::from_iter(3..=5)));
151         assert!(c1.intersect_count(&c2) == 3);
152         assert!(c1.union_count(&c2) == 7);
153     }
154 }
155