rustgym/leetcode/
_295_find_median_from_data_stream.rs

1use std::cmp::Reverse;
2use std::collections::BinaryHeap;
3
4#[derive(Default)]
5struct MedianFinder {
6    lo: BinaryHeap<i32>,
7    hi: BinaryHeap<Reverse<i32>>,
8}
9
10impl MedianFinder {
11    fn new() -> Self {
12        MedianFinder::default()
13    }
14
15    fn add_num(&mut self, num: i32) {
16        self.hi.push(Reverse(num));
17        let smallest = self.hi.pop().unwrap().0;
18        self.lo.push(smallest);
19        if self.lo.len() > self.hi.len() + 1 {
20            self.hi.push(Reverse(self.lo.pop().unwrap()));
21        }
22    }
23
24    fn find_median(&self) -> f64 {
25        if (self.lo.len() + self.hi.len()) % 2 == 0 {
26            (*self.lo.peek().unwrap() + self.hi.peek().unwrap().0) as f64 / 2.0
27        } else {
28            *self.lo.peek().unwrap() as f64
29        }
30    }
31}
32
33#[test]
34fn test() {
35    use assert_approx_eq::assert_approx_eq;
36    let mut obj = MedianFinder::new();
37    obj.add_num(1);
38    obj.add_num(2);
39    assert_approx_eq!(obj.find_median(), 1.5);
40    obj.add_num(3);
41    assert_approx_eq!(obj.find_median(), 2.0);
42}