src/index_heap.rs

   1 use std;
   2 use std::cmp::min;
   3 use std::mem::swap;
   4 use std::ptr;
   5
   6 pub trait Indexing {
   7     fn as_index(&self) -> usize;
   8 }
   9
  10 // 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.
  11 // The elements are sorted by integer keys.
  12 #[derive(Debug)]
  13 pub struct IndexdMinHeap<T: Ord + Indexing> {
  14     positions: Vec<usize>,
  15     data: Vec<T>
  16 }
  17
  18 const TREE_ARITY: usize = 4;
  19 const INVALID_POSITION: usize = std::usize::MAX;
  20
  21 impl<T: Ord + Indexing> IndexdMinHeap<T> {
  22     #[allow(dead_code)]
  23     pub fn new(max_id: usize) -> IndexdMinHeap<T> {
  24         IndexdMinHeap {
  25             positions: vec![INVALID_POSITION; max_id],
  26             data: Vec::new()
  27         }
  28     }
  29
  30     pub fn len(&self) -> usize {
  31         self.data.len()
  32     }
  33
  34     pub fn is_empty(&self) -> bool {
  35         self.len() == 0
  36     }
  37
  38     #[allow(dead_code)]
  39     pub fn contains_index(&self, id: usize) -> bool {
  40         self.positions[id] != INVALID_POSITION
  41     }
  42
  43     #[allow(dead_code)]
  44     pub fn clear(&mut self) {
  45         for element in &self.data {
  46             self.positions[element.as_index()] = INVALID_POSITION;
  47         }
  48         self.data.clear();
  49     }
  50
  51     #[allow(dead_code)]
  52     pub fn peek(&self) -> Option<&T> {
  53         self.data.first()
  54     }
  55
  56     #[allow(dead_code)]
  57     pub fn pop(&mut self) -> Option<T> {
  58         self.data.pop().map(|mut item| {
  59             self.positions[item.as_index()] = INVALID_POSITION;
  60             if !self.is_empty() {
  61                 self.positions[item.as_index()] = 0;
  62                 self.positions[self.data[0].as_index()] = INVALID_POSITION;
  63                 swap(&mut item, &mut self.data[0]);
  64                 self.move_down_in_tree(0);
  65             }
  66             item
  67         })
  68     }
  69
  70     #[allow(dead_code)]
  71     pub fn push(&mut self, element: T) {
  72         assert!(!self.contains_index(element.as_index()));
  73         let insert_position = self.len();
  74         self.positions[element.as_index()] = insert_position;
  75         self.data.push(element);
  76         self.move_up_in_tree(insert_position);
  77     }
  78
  79     // Updates the key of an element if the new key is smaller than the old key.
  80     // Does nothing if the new key is larger.
  81     // Undefined if the element is not part of the queue.
  82     #[allow(dead_code)]
  83     pub fn decrease_key(&mut self, element: T) {
  84         let position = self.positions[element.as_index()];
  85         self.data[position] = element;
  86         self.move_up_in_tree(position);
  87     }
  88
  89     // Updates the key of an element if the new key is larger than the old key.
  90     // Does nothing if the new key is smaller.
  91     // Undefined if the element is not part of the queue.
  92     #[allow(dead_code)]
  93     pub fn increase_key(&mut self, element: T) {
  94         let position = self.positions[element.as_index()];
  95         self.data[position] = element;
  96         self.move_down_in_tree(position);
  97     }
  98
  99     fn move_up_in_tree(&mut self, position: usize) {
 100         unsafe {
 101             let mut position = position;
 102             let mut hole = Hole::new(&mut self.data, position);
 103
 104             while position > 0 {
 105                 let parent = (position - 1) / TREE_ARITY;
 106
 107                 if hole.get(parent) < hole.element() {
 108                     break;
 109                 }
 110
 111                 self.positions[hole.get(parent).as_index()] = position;
 112                 hole.move_to(parent);
 113                 position = parent;
 114             }
 115
 116             self.positions[hole.element().as_index()] = position;
 117         }
 118     }
 119
 120     fn move_down_in_tree(&mut self, position: usize) {
 121         unsafe {
 122             let mut position = position;
 123             let heap_size = self.len();
 124             let mut hole = Hole::new(&mut self.data, position);
 125
 126             loop {
 127                 if let Some(smallest_child) = IndexdMinHeap::<T>::children_index_range(position, heap_size).min_by_key(|&child_index| hole.get(child_index)) {
 128                     if hole.get(smallest_child) >= hole.element() {
 129                         self.positions[hole.element().as_index()] = position;
 130                         return; // no child is smaller
 131                     }
 132
 133                     self.positions[hole.get(smallest_child).as_index()] = position;
 134                     hole.move_to(smallest_child);
 135                     position = smallest_child;
 136                 } else {
 137                     self.positions[hole.element().as_index()] = position;
 138                     return; // no children at all
 139                 }
 140             }
 141         }
 142     }
 143
 144     fn children_index_range(parent_index: usize, heap_size: usize) -> std::ops::Range<usize> {
 145         let first_child = TREE_ARITY * parent_index + 1;
 146         let last_child = min(TREE_ARITY * parent_index + TREE_ARITY + 1, heap_size);
 147         first_child..last_child
 148     }
 149 }
 150
 151
 152 // This is an optimization copied straight from the rust stdlib binary heap
 153 // it allows to avoid always swapping elements pairwise and rather
 154 // move each element only once.
 155
 156
 157 /// Hole represents a hole in a slice i.e. an index without valid value
 158 /// (because it was moved from or duplicated).
 159 /// In drop, `Hole` will restore the slice by filling the hole
 160 /// position with the value that was originally removed.
 161 struct Hole<'a, T: 'a> {
 162     data: &'a mut [T],
 163     /// `elt` is always `Some` from new until drop.
 164     elt: Option<T>,
 165     pos: usize,
 166 }
 167
 168 impl<'a, T> Hole<'a, T> {
 169     /// Create a new Hole at index `pos`.
 170     ///
 171     /// Unsafe because pos must be within the data slice.
 172     #[inline]
 173     unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
 174         debug_assert!(pos < data.len());
 175         let elt = ptr::read(&data[pos]);
 176         Hole {
 177             data,
 178             elt: Some(elt),
 179             pos,
 180         }
 181     }
 182
 183     /// Returns a reference to the element removed.
 184     #[inline]
 185     fn element(&self) -> &T {
 186         self.elt.as_ref().unwrap()
 187     }
 188
 189     /// Returns a reference to the element at `index`.
 190     ///
 191     /// Unsafe because index must be within the data slice and not equal to pos.
 192     #[inline]
 193     unsafe fn get(&self, index: usize) -> &T {
 194         debug_assert!(index != self.pos);
 195         debug_assert!(index < self.data.len());
 196         self.data.get_unchecked(index)
 197     }
 198
 199     /// Move hole to new location
 200     ///
 201     /// Unsafe because index must be within the data slice and not equal to pos.
 202     #[inline]
 203     unsafe fn move_to(&mut self, index: usize) {
 204         debug_assert!(index != self.pos);
 205         debug_assert!(index < self.data.len());
 206         let index_ptr: *const _ = self.data.get_unchecked(index);
 207         let hole_ptr = self.data.get_unchecked_mut(self.pos);
 208         ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
 209         self.pos = index;
 210     }
 211 }
 212
 213 impl<'a, T> Drop for Hole<'a, T> {
 214     #[inline]
 215     fn drop(&mut self) {
 216         // fill the hole again
 217         unsafe {
 218             let pos = self.pos;
 219             ptr::write(self.data.get_unchecked_mut(pos), self.elt.take().unwrap());
 220         }
 221     }
 222 }