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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127
use std::collections::BinaryHeap; use std::cmp::{Ord, Ordering}; #[derive(Debug)] pub struct LimitedBinaryHeap<T> where T: Ord, { binary_heap: BinaryHeap<T>, limit: usize, filled: bool, } impl<T> LimitedBinaryHeap<T> where T: Ord, { pub fn new(limit: usize) -> LimitedBinaryHeap<T> { LimitedBinaryHeap { binary_heap: BinaryHeap::with_capacity(limit as usize), limit, filled: false, } } pub fn insert(&mut self, value: T) { if self.filled { if value.cmp(self.binary_heap.peek().expect("Binary heap empty")) == Ordering::Less { self.binary_heap.pop(); self.binary_heap.push(value); } } else { if self.binary_heap.len() + 1 == self.limit { self.filled = true; } self.binary_heap.push(value); } } pub fn into_sorted_vec(self) -> Vec<T> { self.binary_heap.into_sorted_vec() } } #[cfg(test)] pub mod limited_binary_tests { use super::LimitedBinaryHeap; use std::cmp::{Ord, PartialOrd, Eq, PartialEq, Ordering}; #[derive(Debug)] pub struct RawResponseRecord { pub price: u32, pub product_name: String, } impl Ord for RawResponseRecord { fn cmp(&self, other: &Self) -> Ordering { self.price.cmp(&other.price) } } impl Eq for RawResponseRecord {} impl PartialOrd for RawResponseRecord { fn partial_cmp(&self, other: &Self) -> Option<Ordering> { self.price.partial_cmp(&other.price) } } impl PartialEq for RawResponseRecord { fn eq(&self, other: &Self) -> bool { self.price == other.price } } #[test] fn test_limit() { let mut limited_binary_heap: LimitedBinaryHeap<RawResponseRecord> = LimitedBinaryHeap::new(3); limited_binary_heap.insert(RawResponseRecord { price: 10, product_name: "Seal".to_owned(), }); limited_binary_heap.insert(RawResponseRecord { price: 20, product_name: "Sugar".to_owned(), }); limited_binary_heap.insert(RawResponseRecord { price: 30, product_name: "Pepper".to_owned(), }); limited_binary_heap.insert(RawResponseRecord { price: 5, product_name: "Potato".to_owned(), }); limited_binary_heap.insert(RawResponseRecord { price: 40, product_name: "Milk".to_owned(), }); limited_binary_heap.insert(RawResponseRecord { price: 15, product_name: "Nut".to_owned(), }); assert_eq!( limited_binary_heap.into_sorted_vec(), vec![ RawResponseRecord { price: 5, product_name: "Potato".to_owned(), }, RawResponseRecord { price: 10, product_name: "Seal".to_owned() }, RawResponseRecord { price: 15, product_name: "Nut".to_owned(), }, ] ); } }