1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82
use ahash::AHashMap; #[derive(Default)] pub struct Counter { counts: AHashMap<String, u64>, top: AHashMap<String, u64>, threshold: u64, num: usize, } impl Counter { pub fn new(num: Option<usize>) -> Self { Self { num: num.unwrap_or(0), threshold: num.map(|_| 0).unwrap_or(u64::MAX), ..Default::default() } } pub fn add(&mut self, key: &str, added: u64) { let count = match self.counts.get_mut(key) { Some(count) => { *count += added; *count } None => { self.counts.insert(key.to_owned(), added); added } }; if count < self.threshold { return; } if let Some(t) = self.top.get_mut(key) { *t = count; return; } self.top.insert(key.to_owned(), count); if self.top.len() < self.num * 2 { return; } let mut top_values = self.top.values().collect::<Vec<_>>(); top_values.sort_unstable(); let threshold = *top_values[self.num as usize - 1]; self.threshold = threshold; self.top.retain(|_, v| *v > threshold); } pub fn top(&self) -> Vec<KeyCount> { let mut top = Vec::with_capacity(self.num); for (key, &count) in &self.top { top.push(KeyCount { count, key: key.into(), }); } top.sort_unstable(); top.reverse(); if self.num > 0 { top.truncate(self.num); } top } pub fn merge(mut self, r: Counter) -> Self { for (key, count) in r.counts.iter() { self.add(key, *count); } self } } #[derive(PartialEq, PartialOrd, Eq, Ord)] pub struct KeyCount { pub count: u64, pub key: String, }