Line data    Source code

       1              : //! A top-k heap.
       2              : //!
       3              : //! The heap is a min heap over at most `capacity` entries: a new entry replaces the minimum if the
       4              : //! minimum is less. The order should be total, then the top k and their order are well-defined.
       5              : 
       6              : use std::cmp::Ordering;
       7              : 
       8              : /// A min heap that keeps the `capacity` greatest entries
       9              : pub struct TopKHeap<T, C: Fn(&T, &T) -> Ordering> {
      10              :     /// The entries, a min heap once full and sorted in descending order when finished
      11              :     entries: Vec<T>,
      12              :     /// The maximum number of entries
      13              :     capacity: usize,
      14              :     /// The total order of the entries
      15              :     cmp: C,
      16              : }
      17              : 
      18              : impl<T, C: Fn(&T, &T) -> Ordering> TopKHeap<T, C> {
      19          992 :     pub fn new(capacity: usize, cmp: C) -> Self {
      20          992 :         Self { entries: Vec::with_capacity(capacity), capacity, cmp }
      21          992 :     }
      22              : 
      23        85830 :     fn less(&self, l: usize, r: usize) -> bool {
      24        85830 :         (self.cmp)(&self.entries[l], &self.entries[r]) == Ordering::Less
      25        85830 :     }
      26              : 
      27              :     /// Move an entry down until the first `len` entries are a min heap below it
      28        15165 :     fn sift_down(&mut self, mut i: usize, len: usize) {
      29              :         loop {
      30        55181 :             let l = 2 * i + 1;
      31        55181 :             let r = 2 * i + 2;
      32        55181 :             let mut smallest = i;
      33        55181 :             if l < len && self.less(l, smallest) {
      34        37294 :                 smallest = l;
      35        37503 :             }
      36        55181 :             if r < len && self.less(r, smallest) {
      37        17770 :                 smallest = r;
      38        37411 :             }
      39        55181 :             if smallest == i {
      40        15165 :                 break;
      41        40016 :             }
      42        40016 :             self.entries.swap(i, smallest);
      43        40016 :             i = smallest;
      44              :         }
      45        15165 :     }
      46              : 
      47              :     /// Make the entries a min heap
      48          815 :     fn heapify(&mut self) {
      49          815 :         let len = self.entries.len();
      50         2918 :         for i in (0..len / 2).rev() {
      51         2768 :             self.sift_down(i, len);
      52         2768 :         }
      53          815 :     }
      54              : 
      55        79026 :     pub fn insert(&mut self, value: T) {
      56        79026 :         if self.entries.len() < self.capacity {
      57         6056 :             self.entries.push(value);
      58         6056 :             if self.entries.len() == self.capacity {
      59          221 :                 self.heapify();
      60         5835 :             }
      61        72970 :         } else if self.capacity > 0 && (self.cmp)(&self.entries[0], &value) == Ordering::Less {
      62         7080 :             self.entries[0] = value;
      63         7080 :             self.sift_down(0, self.capacity);
      64        65890 :         }
      65        79026 :     }
      66              : 
      67          361 :     pub fn clear(&mut self) {
      68          361 :         self.entries.clear();
      69          361 :     }
      70              : 
      71              :     /// Sort the entries in descending order: heap sort, moving the minimum to the end
      72          815 :     pub fn finish(&mut self) -> &mut Vec<T> {
      73          815 :         if self.entries.len() < self.capacity {
      74          594 :             self.heapify();
      75          597 :         }
      76         5319 :         for end in (1..self.entries.len()).rev() {
      77         5317 :             self.entries.swap(0, end);
      78         5317 :             self.sift_down(0, end);
      79         5317 :         }
      80          815 :         &mut self.entries
      81          815 :     }
      82              : 
      83            0 :     pub fn entries(&self) -> &[T] {
      84            0 :         &self.entries
      85            0 :     }
      86              : 
      87          446 :     pub fn take_entries(&mut self) -> Vec<T> {
      88          446 :         std::mem::take(&mut self.entries)
      89          446 :     }
      90              : }