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          952 :     pub fn new(capacity: usize, cmp: C) -> Self {
      20          952 :         Self { entries: Vec::with_capacity(capacity), capacity, cmp }
      21          952 :     }
      22              : 
      23        86987 :     fn less(&self, l: usize, r: usize) -> bool {
      24        86987 :         (self.cmp)(&self.entries[l], &self.entries[r]) == Ordering::Less
      25        86987 :     }
      26              : 
      27              :     /// Move an entry down until the first `len` entries are a min heap below it
      28        15108 :     fn sift_down(&mut self, mut i: usize, len: usize) {
      29              :         loop {
      30        55747 :             let l = 2 * i + 1;
      31        55747 :             let r = 2 * i + 2;
      32        55747 :             let mut smallest = i;
      33        55747 :             if l < len && self.less(l, smallest) {
      34        37782 :                 smallest = l;
      35        37972 :             }
      36        55747 :             if r < len && self.less(r, smallest) {
      37        18097 :                 smallest = r;
      38        37650 :             }
      39        55747 :             if smallest == i {
      40        15108 :                 break;
      41        40639 :             }
      42        40639 :             self.entries.swap(i, smallest);
      43        40639 :             i = smallest;
      44              :         }
      45        15108 :     }
      46              : 
      47              :     /// Make the entries a min heap
      48          805 :     fn heapify(&mut self) {
      49          805 :         let len = self.entries.len();
      50         2866 :         for i in (0..len / 2).rev() {
      51         2724 :             self.sift_down(i, len);
      52         2724 :         }
      53          805 :     }
      54              : 
      55        75790 :     pub fn insert(&mut self, value: T) {
      56        75790 :         if self.entries.len() < self.capacity {
      57         6027 :             self.entries.push(value);
      58         6027 :             if self.entries.len() == self.capacity {
      59          210 :                 self.heapify();
      60         5817 :             }
      61        69763 :         } else if self.capacity > 0 && (self.cmp)(&self.entries[0], &value) == Ordering::Less {
      62         7087 :             self.entries[0] = value;
      63         7087 :             self.sift_down(0, self.capacity);
      64        62676 :         }
      65        75790 :     }
      66              : 
      67          352 :     pub fn clear(&mut self) {
      68          352 :         self.entries.clear();
      69          352 :     }
      70              : 
      71              :     /// Sort the entries in descending order: heap sort, moving the minimum to the end
      72          805 :     pub fn finish(&mut self) -> &mut Vec<T> {
      73          805 :         if self.entries.len() < self.capacity {
      74          595 :             self.heapify();
      75          598 :         }
      76         5299 :         for end in (1..self.entries.len()).rev() {
      77         5297 :             self.entries.swap(0, end);
      78         5297 :             self.sift_down(0, end);
      79         5297 :         }
      80          805 :         &mut self.entries
      81          805 :     }
      82              : 
      83            0 :     pub fn entries(&self) -> &[T] {
      84            0 :         &self.entries
      85            0 :     }
      86              : 
      87          445 :     pub fn take_entries(&mut self) -> Vec<T> {
      88          445 :         std::mem::take(&mut self.entries)
      89          445 :     }
      90              : }