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