src/index_heap.rs

   1 //! A priority queue implemented with a 4-ary heap.
   2 //!
   3 //! Insertion and popping the minimal element have `O(log n)` time complexity.
   4 //! Checking the minimal element is `O(1)`. Keys of elements in the heap can
   5 //! also be increased or decreased
   6 //!
   7 //! # Examples
   8 //!
   9 //! ```
  10 //! use stud_rust_base::index_heap::{Indexing, IndexdMinHeap};
  11 //!
  12 //! #[derive(Copy, Clone, Eq, PartialEq, Debug, Ord, PartialOrd)]
  13 //! pub struct State {
  14 //!     pub distance: usize,
  15 //!     pub node: usize,
  16 //! }
  17 //!
  18 //!
  19 //! // The `Indexing` traits needs to be implemented as well, so we can find elements to decrease their key.
  20 //! impl Indexing for State {
  21 //!     fn as_index(&self) -> usize {
  22 //!         self.node as usize
  23 //!     }
  24 //! }
  25 //!
  26 //! let mut heap = IndexdMinHeap::new(3);
  27 //! heap.push(State { node: 0, distance: 42 });
  28 //! heap.push(State { node: 1, distance: 23 });
  29 //! heap.push(State { node: 2, distance: 50000 });
  30 //! assert_eq!(heap.peek().cloned(), Some(State { node: 1, distance: 23 }));
  31 //! heap.decrease_key(State { node: 0, distance: 1 });
  32 //! assert_eq!(heap.pop(), Some(State { node: 0, distance: 1 }));
  33 //!
  34 //! ```
  35
  36 use std;
  37 use std::cmp::min;
  38 use std::mem::swap;
  39 use std::ptr;
  40
  41 /// A trait to map elements in a heap to a unique index.
  42 /// The element type of the `IndexdMinHeap` has to implement this trait.
  43 pub trait Indexing {
  44     /// This method has to map a heap element to a unique `usize` index.
  45     fn as_index(&self) -> usize;
  46 }
  47
  48 /// A priority queue where the elements are IDs from 0 to id_count-1 where id_count is a number that is set in the constructor.
  49 /// The elements are sorted ascending by the ordering defined by the `Ord` trait.
  50 /// The interface mirros the standard library BinaryHeap (except for the reversed order).
  51 /// Only the methods necessary for dijkstras algorithm are implemented.
  52 /// In addition, `increase_key` and `decrease_key` methods are available.
  53 #[derive(Debug)]
  54 pub struct IndexdMinHeap<T: Ord + Indexing> {
  55     positions: Vec<usize>,
  56     data: Vec<T>,
  57 }
  58
  59 const TREE_ARITY: usize = 4;
  60 const INVALID_POSITION: usize = std::usize::MAX;
  61
  62 impl<T: Ord + Indexing> IndexdMinHeap<T> {
  63     /// Creates an empty `IndexdMinHeap` as a min-heap.
  64     /// The indices (as defined by the `Indexing` trait) of all inserted elements
  65     /// will have to be between in `[0, max_index)`
  66     pub fn new(max_index: usize) -> IndexdMinHeap<T> {
  67         IndexdMinHeap {
  68             positions: vec![INVALID_POSITION; max_index],
  69             data: Vec::new(),
  70         }
  71     }
  72
  73     /// Returns the length of the binary heap.
  74     pub fn len(&self) -> usize {
  75         self.data.len()
  76     }
  77
  78     /// Checks if the binary heap is empty.
  79     pub fn is_empty(&self) -> bool {
  80         self.len() == 0
  81     }
  82
  83     /// Checks if the heap already contains an element mapped to the given index
  84     pub fn contains_index(&self, id: usize) -> bool {
  85         self.positions[id] != INVALID_POSITION
  86     }
  87
  88     /// Drops all items from the heap.
  89     pub fn clear(&mut self) {
  90         for element in &self.data {
  91             self.positions[element.as_index()] = INVALID_POSITION;
  92         }
  93         self.data.clear();
  94     }
  95
  96     /// Returns a reference to the smallest item in the heap, or None if it is empty.
  97     pub fn peek(&self) -> Option<&T> {
  98         self.data.first()
  99     }
 100
 101     /// Removes the greatest item from the binary heap and returns it, or None if it is empty.
 102     pub fn pop(&mut self) -> Option<T> {
 103         self.data.pop().map(|mut item| {
 104             self.positions[item.as_index()] = INVALID_POSITION;
 105             if !self.is_empty() {
 106                 self.positions[item.as_index()] = 0;
 107                 self.positions[self.data[0].as_index()] = INVALID_POSITION;
 108                 swap(&mut item, &mut self.data[0]);
 109                 self.move_down_in_tree(0);
 110             }
 111             item
 112         })
 113     }
 114
 115     /// Pushes an item onto the binary heap.
 116     /// Panics if an element with the same index already exists.
 117     pub fn push(&mut self, element: T) {
 118         assert!(!self.contains_index(element.as_index()));
 119         let insert_position = self.len();
 120         self.positions[element.as_index()] = insert_position;
 121         self.data.push(element);
 122         self.move_up_in_tree(insert_position);
 123     }
 124
 125     /// Updates the key of an element if the new key is smaller than the old key.
 126     /// Panics if the element is not part of the queue or if the new key is larger.
 127     pub fn decrease_key(&mut self, element: T) {
 128         let position = self.positions[element.as_index()];
 129         assert!(element <= self.data[position]);
 130         self.data[position] = element;
 131         self.move_up_in_tree(position);
 132     }
 133
 134     /// Updates the key of an element if the new key is larger than the old key.
 135     /// Panics if the element is not part of the queue or if the new key is smaller.
 136     pub fn increase_key(&mut self, element: T) {
 137         let position = self.positions[element.as_index()];
 138         assert!(element >= self.data[position]);
 139         self.data[position] = element;
 140         self.move_down_in_tree(position);
 141     }
 142
 143     fn move_up_in_tree(&mut self, position: usize) {
 144         unsafe {
 145             let mut position = position;
 146             let mut hole = Hole::new(&mut self.data, position);
 147
 148             while position > 0 {
 149                 let parent = (position - 1) / TREE_ARITY;
 150
 151                 if hole.get(parent) < hole.element() {
 152                     break;
 153                 }
 154
 155                 self.positions[hole.get(parent).as_index()] = position;
 156                 hole.move_to(parent);
 157                 position = parent;
 158             }
 159
 160             self.positions[hole.element().as_index()] = position;
 161         }
 162     }
 163
 164     fn move_down_in_tree(&mut self, position: usize) {
 165         unsafe {
 166             let mut position = position;
 167             let heap_size = self.len();
 168             let mut hole = Hole::new(&mut self.data, position);
 169
 170             loop {
 171                 if let Some(smallest_child) = IndexdMinHeap::<T>::children_index_range(position, heap_size).min_by_key(|&child_index| hole.get(child_index)) {
 172                     if hole.get(smallest_child) >= hole.element() {
 173                         self.positions[hole.element().as_index()] = position;
 174                         return; // no child is smaller
 175                     }
 176
 177                     self.positions[hole.get(smallest_child).as_index()] = position;
 178                     hole.move_to(smallest_child);
 179                     position = smallest_child;
 180                 } else {
 181                     self.positions[hole.element().as_index()] = position;
 182                     return; // no children at all
 183                 }
 184             }
 185         }
 186     }
 187
 188     fn children_index_range(parent_index: usize, heap_size: usize) -> std::ops::Range<usize> {
 189         let first_child = TREE_ARITY * parent_index + 1;
 190         let last_child = min(TREE_ARITY * parent_index + TREE_ARITY + 1, heap_size);
 191         first_child..last_child
 192     }
 193 }
 194
 195 // This is an optimization copied straight from the rust stdlib binary heap
 196 // it allows to avoid always swapping elements pairwise and rather
 197 // move each element only once.
 198
 199 /// Hole represents a hole in a slice i.e. an index without valid value
 200 /// (because it was moved from or duplicated).
 201 /// In drop, `Hole` will restore the slice by filling the hole
 202 /// position with the value that was originally removed.
 203 struct Hole<'a, T: 'a> {
 204     data: &'a mut [T],
 205     /// `elt` is always `Some` from new until drop.
 206     elt: Option<T>,
 207     pos: usize,
 208 }
 209
 210 impl<'a, T> Hole<'a, T> {
 211     /// Create a new Hole at index `pos`.
 212     ///
 213     /// Unsafe because pos must be within the data slice.
 214     #[inline]
 215     unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
 216         debug_assert!(pos < data.len());
 217         let elt = ptr::read(&data[pos]);
 218         Hole { data, elt: Some(elt), pos }
 219     }
 220
 221     /// Returns a reference to the element removed.
 222     #[inline]
 223     fn element(&self) -> &T {
 224         self.elt.as_ref().unwrap()
 225     }
 226
 227     /// Returns a reference to the element at `index`.
 228     ///
 229     /// Unsafe because index must be within the data slice and not equal to pos.
 230     #[inline]
 231     unsafe fn get(&self, index: usize) -> &T {
 232         debug_assert!(index != self.pos);
 233         debug_assert!(index < self.data.len());
 234         self.data.get_unchecked(index)
 235     }
 236
 237     /// Move hole to new location
 238     ///
 239     /// Unsafe because index must be within the data slice and not equal to pos.
 240     #[inline]
 241     unsafe fn move_to(&mut self, index: usize) {
 242         debug_assert!(index != self.pos);
 243         debug_assert!(index < self.data.len());
 244         let index_ptr: *const _ = self.data.get_unchecked(index);
 245         let hole_ptr = self.data.get_unchecked_mut(self.pos);
 246         ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
 247         self.pos = index;
 248     }
 249 }
 250
 251 impl<'a, T> Drop for Hole<'a, T> {
 252     #[inline]
 253     fn drop(&mut self) {
 254         // fill the hole again
 255         unsafe {
 256             let pos = self.pos;
 257             ptr::write(self.data.get_unchecked_mut(pos), self.elt.take().unwrap());
 258         }
 259     }
 260 }