competitive/data_structure/
counter.rs

1use std::{
2    borrow::Borrow,
3    collections::{BTreeMap, HashMap, btree_map, hash_map},
4    fmt::{self, Debug},
5    hash::Hash,
6    iter::{Extend, FromIterator},
7    ops::RangeBounds,
8};
9
10#[derive(Clone)]
11pub struct HashCounter<T> {
12    map: HashMap<T, usize>,
13}
14impl<T> Debug for HashCounter<T>
15where
16    T: Debug,
17{
18    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
19        f.debug_map().entries(self.iter()).finish()
20    }
21}
22impl<T> Default for HashCounter<T>
23where
24    T: Eq + Hash,
25{
26    fn default() -> Self {
27        Self {
28            map: HashMap::default(),
29        }
30    }
31}
32impl<T> HashCounter<T> {
33    pub fn len(&self) -> usize {
34        self.map.len()
35    }
36    pub fn is_empty(&self) -> bool {
37        self.map.is_empty()
38    }
39    pub fn clear(&mut self) {
40        self.map.clear();
41    }
42    pub fn keys(&self) -> hash_map::Keys<'_, T, usize> {
43        self.map.keys()
44    }
45    pub fn values(&self) -> hash_map::Values<'_, T, usize> {
46        self.map.values()
47    }
48    pub fn iter(&self) -> hash_map::Iter<'_, T, usize> {
49        self.map.iter()
50    }
51    pub fn drain(&mut self) -> hash_map::Drain<'_, T, usize> {
52        self.map.drain()
53    }
54}
55impl<T> HashCounter<T>
56where
57    T: Eq + Hash,
58{
59    pub fn new() -> Self {
60        Self::default()
61    }
62    pub fn with_capacity(capacity: usize) -> Self {
63        Self {
64            map: HashMap::with_capacity(capacity),
65        }
66    }
67    pub fn get(&self, item: &T) -> usize {
68        self.map.get(item).cloned().unwrap_or_default()
69    }
70    pub fn add(&mut self, item: T) {
71        self.add_count(item, 1)
72    }
73    pub fn add_count(&mut self, item: T, count: usize) {
74        *self.map.entry(item).or_default() += count;
75    }
76    pub fn remove(&mut self, item: &T) -> bool {
77        self.remove_count(item, 1) == 1
78    }
79    pub fn remove_count(&mut self, item: &T, count: usize) -> usize {
80        if let Some(cnt) = self.map.get_mut(item) {
81            if *cnt <= count {
82                let cnt = *cnt;
83                self.map.remove(item);
84                cnt
85            } else {
86                *cnt -= count;
87                count
88            }
89        } else {
90            0
91        }
92    }
93    pub fn append(&mut self, other: &mut Self) {
94        if self.map.len() < other.map.len() {
95            std::mem::swap(self, other);
96        }
97        self.extend(other.map.drain());
98    }
99}
100impl<T> Extend<T> for HashCounter<T>
101where
102    T: Eq + Hash,
103{
104    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
105        for item in iter {
106            self.add(item)
107        }
108    }
109}
110impl<T> Extend<(T, usize)> for HashCounter<T>
111where
112    T: Eq + Hash,
113{
114    fn extend<I: IntoIterator<Item = (T, usize)>>(&mut self, iter: I) {
115        for (item, count) in iter {
116            self.add_count(item, count)
117        }
118    }
119}
120impl<T> FromIterator<T> for HashCounter<T>
121where
122    T: Eq + Hash,
123{
124    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
125        let mut map = Self::default();
126        map.extend(iter);
127        map
128    }
129}
130impl<T> FromIterator<(T, usize)> for HashCounter<T>
131where
132    T: Eq + Hash,
133{
134    fn from_iter<I: IntoIterator<Item = (T, usize)>>(iter: I) -> Self {
135        let mut map = Self::default();
136        map.extend(iter);
137        map
138    }
139}
140
141#[derive(Clone)]
142pub struct BTreeCounter<T> {
143    map: BTreeMap<T, usize>,
144}
145impl<T> Debug for BTreeCounter<T>
146where
147    T: Debug,
148{
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.debug_map().entries(self.iter()).finish()
151    }
152}
153impl<T> Default for BTreeCounter<T>
154where
155    T: Ord,
156{
157    fn default() -> Self {
158        Self {
159            map: BTreeMap::default(),
160        }
161    }
162}
163impl<T> BTreeCounter<T> {
164    pub fn len(&self) -> usize {
165        self.map.len()
166    }
167    pub fn is_empty(&self) -> bool {
168        self.map.is_empty()
169    }
170    pub fn keys(&self) -> btree_map::Keys<'_, T, usize> {
171        self.map.keys()
172    }
173    pub fn values(&self) -> btree_map::Values<'_, T, usize> {
174        self.map.values()
175    }
176    pub fn iter(&self) -> btree_map::Iter<'_, T, usize> {
177        self.map.iter()
178    }
179    pub fn range<Q, R>(&self, range: R) -> btree_map::Range<'_, T, usize>
180    where
181        Q: Ord,
182        R: RangeBounds<Q>,
183        T: Borrow<Q> + Ord,
184    {
185        self.map.range(range)
186    }
187}
188impl<T> BTreeCounter<T>
189where
190    T: Ord,
191{
192    pub fn new() -> Self {
193        Self::default()
194    }
195    pub fn clear(&mut self) {
196        self.map.clear();
197    }
198    pub fn get(&self, item: &T) -> usize {
199        self.map.get(item).cloned().unwrap_or_default()
200    }
201    pub fn add(&mut self, item: T) {
202        self.add_count(item, 1)
203    }
204    pub fn add_count(&mut self, item: T, count: usize) {
205        *self.map.entry(item).or_default() += count;
206    }
207    pub fn remove(&mut self, item: &T) -> bool {
208        self.remove_count(item, 1) == 1
209    }
210    pub fn remove_count(&mut self, item: &T, count: usize) -> usize {
211        if let Some(cnt) = self.map.get_mut(item) {
212            if *cnt <= count {
213                let cnt = *cnt;
214                self.map.remove(item);
215                cnt
216            } else {
217                *cnt -= count;
218                count
219            }
220        } else {
221            0
222        }
223    }
224}
225impl<T> Extend<T> for BTreeCounter<T>
226where
227    T: Ord,
228{
229    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
230        for item in iter {
231            self.add(item)
232        }
233    }
234}
235impl<T> Extend<(T, usize)> for BTreeCounter<T>
236where
237    T: Ord,
238{
239    fn extend<I: IntoIterator<Item = (T, usize)>>(&mut self, iter: I) {
240        for (item, count) in iter {
241            self.add_count(item, count)
242        }
243    }
244}
245impl<T> FromIterator<T> for BTreeCounter<T>
246where
247    T: Ord,
248{
249    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
250        let mut map = Self::default();
251        map.extend(iter);
252        map
253    }
254}
255impl<T> FromIterator<(T, usize)> for BTreeCounter<T>
256where
257    T: Ord,
258{
259    fn from_iter<I: IntoIterator<Item = (T, usize)>>(iter: I) -> Self {
260        let mut map = Self::default();
261        map.extend(iter);
262        map
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use crate::tools::Xorshift;
270
271    #[test]
272    fn test_hash_counter() {
273        let mut rng = Xorshift::default();
274        const N: usize = 100_000;
275        const Q: usize = 100_000;
276        let mut cnt = HashCounter::<usize>::new();
277        let mut arr = vec![0usize; N];
278        for _ in 0..Q {
279            let x = rng.random(..N);
280            let c = rng.random(..N);
281            assert_eq!(cnt.get(&x), arr[x]);
282            match rng.random(0..4) {
283                0 => {
284                    cnt.add(x);
285                    arr[x] += 1;
286                }
287                1 => {
288                    assert_eq!(cnt.remove(&x), arr[x] > 0);
289                    if arr[x] > 0 {
290                        arr[x] -= 1;
291                    }
292                }
293                2 => {
294                    cnt.add_count(x, c);
295                    arr[x] += c;
296                }
297                _ => {
298                    assert_eq!(cnt.remove_count(&x, c), arr[x].min(c));
299                    arr[x] -= arr[x].min(c);
300                }
301            }
302            assert_eq!(cnt.get(&x), arr[x]);
303        }
304    }
305
306    #[test]
307    fn test_btree_counter() {
308        let mut rng = Xorshift::default();
309        const N: usize = 100_000;
310        const Q: usize = 100_000;
311        let mut cnt = BTreeCounter::<usize>::new();
312        let mut arr = vec![0usize; N];
313        for _ in 0..Q {
314            let x = rng.random(..N);
315            let c = rng.random(..N);
316            assert_eq!(cnt.get(&x), arr[x]);
317            match rng.random(0..4) {
318                0 => {
319                    cnt.add(x);
320                    arr[x] += 1;
321                }
322                1 => {
323                    assert_eq!(cnt.remove(&x), arr[x] > 0);
324                    if arr[x] > 0 {
325                        arr[x] -= 1;
326                    }
327                }
328                2 => {
329                    cnt.add_count(x, c);
330                    arr[x] += c;
331                }
332                _ => {
333                    assert_eq!(cnt.remove_count(&x, c), arr[x].min(c));
334                    arr[x] -= arr[x].min(c);
335                }
336            }
337            assert_eq!(cnt.get(&x), arr[x]);
338        }
339    }
340}
competitive/data_structure/counter.rs

competitive/data_structure/
counter.rs
