tycho_util/num/
median.rs

1use std::cmp::Reverse;
2use std::collections::BinaryHeap;
3
4#[derive(Default)]
5pub struct StreamingUnsignedMedian {
6    lo: BinaryHeap<u128>,          // max-heap
7    hi: BinaryHeap<Reverse<u128>>, // min-heap
8}
9
10impl StreamingUnsignedMedian {
11    pub fn insert<V: Into<u128>>(&mut self, v: V) {
12        let v = v.into();
13        // 1) put to hi, move the minimal to lo
14        self.hi.push(Reverse(v));
15        if let Some(Reverse(v)) = self.hi.pop() {
16            self.lo.push(v);
17        }
18        // 2) rebalance
19        if self.lo.len() > self.hi.len() + 1 {
20            let v = self.lo.pop().unwrap();
21            self.hi.push(Reverse(v));
22        }
23    }
24
25    pub fn get_median(&self) -> u128 {
26        self.get_median_checked()
27            .expect("should check first with `median_checked`")
28    }
29
30    pub fn get_median_checked(&self) -> Option<u128> {
31        if self.lo.is_empty() {
32            return None;
33        }
34        if self.lo.len() > self.hi.len() {
35            Some(*self.lo.peek().unwrap())
36        } else {
37            let a = *self.lo.peek().unwrap();
38            let b = self.hi.peek().unwrap().0;
39            Some(a.saturating_add(b).saturating_div(2))
40        }
41    }
42}
43
44pub struct VecOfStreamingUnsignedMedian {
45    idx: usize,
46    inner: Vec<StreamingUnsignedMedian>,
47}
48
49impl VecOfStreamingUnsignedMedian {
50    pub fn new(capacity: usize) -> Self {
51        let mut vec = Vec::with_capacity(capacity);
52        for _ in 0..capacity {
53            vec.push(StreamingUnsignedMedian::default());
54        }
55        Self { idx: 0, inner: vec }
56    }
57
58    pub fn accum<V: Into<u128>>(&mut self, idx: usize, v: V) {
59        self.insert(idx, v);
60    }
61
62    pub fn accum_next<V: Into<u128>>(&mut self, v: V) {
63        self.insert_next(v);
64    }
65
66    pub fn insert<V: Into<u128>>(&mut self, idx: usize, v: V) {
67        self.idx = idx;
68        self.insert_curr(v);
69    }
70
71    pub fn insert_next<V: Into<u128>>(&mut self, v: V) {
72        self.idx += 1;
73        self.insert_curr(v);
74    }
75
76    fn insert_curr<V: Into<u128>>(&mut self, v: V) {
77        self.inner[self.idx].insert(v);
78    }
79
80    pub fn get_avg(&mut self, idx: usize) -> u128 {
81        self.get_median(idx)
82    }
83
84    pub fn get_avg_checked(&mut self, idx: usize) -> Option<u128> {
85        self.get_median_checked(idx)
86    }
87
88    pub fn get_avg_next(&mut self) -> u128 {
89        self.get_median_next()
90    }
91
92    pub fn get_avg_next_checked(&mut self) -> Option<u128> {
93        self.get_median_next_checked()
94    }
95
96    pub fn get_median(&mut self, idx: usize) -> u128 {
97        self.get_median_checked(idx)
98            .expect("should check first with `get_median_checked`")
99    }
100
101    pub fn get_median_checked(&mut self, idx: usize) -> Option<u128> {
102        self.idx = idx;
103        self.get_median_curr()
104    }
105
106    pub fn get_median_next(&mut self) -> u128 {
107        self.get_median_next_checked()
108            .expect("should check first with `get_median_checked`")
109    }
110
111    pub fn get_median_next_checked(&mut self) -> Option<u128> {
112        self.idx += 1;
113        self.get_median_curr()
114    }
115
116    fn get_median_curr(&self) -> Option<u128> {
117        self.inner[self.idx].get_median_checked()
118    }
119}
tycho_util/num/median.rs

tycho_util/num/
median.rs
