src/libcollections/priority_queue.rs

   1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! A priority queue implemented with a binary heap
  12
  13 #![allow(missing_doc)]
  14
  15 use std::clone::Clone;
  16 use std::mem::{move_val_init, init, replace, swap};
  17 use std::slice;
  18
  19 /// A priority queue implemented with a binary heap
  20 #[deriving(Clone)]
  21 pub struct PriorityQueue<T> {
  22     priv data: ~[T],
  23 }
  24
  25 impl<T:Ord> Container for PriorityQueue<T> {
  26     /// Returns the length of the queue
  27     fn len(&self) -> uint { self.data.len() }
  28 }
  29
  30 impl<T:Ord> Mutable for PriorityQueue<T> {
  31     /// Drop all items from the queue
  32     fn clear(&mut self) { self.data.truncate(0) }
  33 }
  34
  35 impl<T:Ord> PriorityQueue<T> {
  36     /// An iterator visiting all values in underlying vector, in
  37     /// arbitrary order.
  38     pub fn iter<'a>(&'a self) -> Items<'a, T> {
  39         Items { iter: self.data.iter() }
  40     }
  41
  42     /// Returns the greatest item in the queue - fails if empty
  43     pub fn top<'a>(&'a self) -> &'a T { &self.data[0] }
  44
  45     /// Returns the greatest item in the queue - None if empty
  46     pub fn maybe_top<'a>(&'a self) -> Option<&'a T> {
  47         if self.is_empty() { None } else { Some(self.top()) }
  48     }
  49
  50     /// Returns the number of elements the queue can hold without reallocating
  51     pub fn capacity(&self) -> uint { self.data.capacity() }
  52
  53     /// Reserve capacity for exactly n elements in the PriorityQueue.
  54     /// Do nothing if the capacity is already sufficient.
  55     pub fn reserve_exact(&mut self, n: uint) { self.data.reserve_exact(n) }
  56
  57     /// Reserve capacity for at least n elements in the PriorityQueue.
  58     /// Do nothing if the capacity is already sufficient.
  59     pub fn reserve(&mut self, n: uint) {
  60         self.data.reserve(n)
  61     }
  62
  63     /// Pop the greatest item from the queue - fails if empty
  64     pub fn pop(&mut self) -> T {
  65         let mut item = self.data.pop().unwrap();
  66         if !self.is_empty() {
  67             swap(&mut item, &mut self.data[0]);
  68             self.siftdown(0);
  69         }
  70         item
  71     }
  72
  73     /// Pop the greatest item from the queue - None if empty
  74     pub fn maybe_pop(&mut self) -> Option<T> {
  75         if self.is_empty() { None } else { Some(self.pop()) }
  76     }
  77
  78     /// Push an item onto the queue
  79     pub fn push(&mut self, item: T) {
  80         self.data.push(item);
  81         let new_len = self.len() - 1;
  82         self.siftup(0, new_len);
  83     }
  84
  85     /// Optimized version of a push followed by a pop
  86     pub fn push_pop(&mut self, mut item: T) -> T {
  87         if !self.is_empty() && self.data[0] > item {
  88             swap(&mut item, &mut self.data[0]);
  89             self.siftdown(0);
  90         }
  91         item
  92     }
  93
  94     /// Optimized version of a pop followed by a push - fails if empty
  95     pub fn replace(&mut self, mut item: T) -> T {
  96         swap(&mut item, &mut self.data[0]);
  97         self.siftdown(0);
  98         item
  99     }
 100
 101     /// Consume the PriorityQueue and return the underlying vector
 102     pub fn to_vec(self) -> ~[T] { let PriorityQueue{data: v} = self; v }
 103
 104     /// Consume the PriorityQueue and return a vector in sorted
 105     /// (ascending) order
 106     pub fn to_sorted_vec(self) -> ~[T] {
 107         let mut q = self;
 108         let mut end = q.len();
 109         while end > 1 {
 110             end -= 1;
 111             q.data.swap(0, end);
 112             q.siftdown_range(0, end)
 113         }
 114         q.to_vec()
 115     }
 116
 117     /// Create an empty PriorityQueue
 118     pub fn new() -> PriorityQueue<T> { PriorityQueue{data: ~[],} }
 119
 120     /// Create a PriorityQueue from a vector (heapify)
 121     pub fn from_vec(xs: ~[T]) -> PriorityQueue<T> {
 122         let mut q = PriorityQueue{data: xs,};
 123         let mut n = q.len() / 2;
 124         while n > 0 {
 125             n -= 1;
 126             q.siftdown(n)
 127         }
 128         q
 129     }
 130
 131     // The implementations of siftup and siftdown use unsafe blocks in
 132     // order to move an element out of the vector (leaving behind a
 133     // zeroed element), shift along the others and move it back into the
 134     // vector over the junk element.  This reduces the constant factor
 135     // compared to using swaps, which involves twice as many moves.
 136     fn siftup(&mut self, start: uint, mut pos: uint) {
 137         unsafe {
 138             let new = replace(&mut self.data[pos], init());
 139
 140             while pos > start {
 141                 let parent = (pos - 1) >> 1;
 142                 if new > self.data[parent] {
 143                     let x = replace(&mut self.data[parent], init());
 144                     move_val_init(&mut self.data[pos], x);
 145                     pos = parent;
 146                     continue
 147                 }
 148                 break
 149             }
 150             move_val_init(&mut self.data[pos], new);
 151         }
 152     }
 153
 154     fn siftdown_range(&mut self, mut pos: uint, end: uint) {
 155         unsafe {
 156             let start = pos;
 157             let new = replace(&mut self.data[pos], init());
 158
 159             let mut child = 2 * pos + 1;
 160             while child < end {
 161                 let right = child + 1;
 162                 if right < end && !(self.data[child] > self.data[right]) {
 163                     child = right;
 164                 }
 165                 let x = replace(&mut self.data[child], init());
 166                 move_val_init(&mut self.data[pos], x);
 167                 pos = child;
 168                 child = 2 * pos + 1;
 169             }
 170
 171             move_val_init(&mut self.data[pos], new);
 172             self.siftup(start, pos);
 173         }
 174     }
 175
 176     fn siftdown(&mut self, pos: uint) {
 177         let len = self.len();
 178         self.siftdown_range(pos, len);
 179     }
 180 }
 181
 182 /// PriorityQueue iterator
 183 pub struct Items <'a, T> {
 184     priv iter: slice::Items<'a, T>,
 185 }
 186
 187 impl<'a, T> Iterator<&'a T> for Items<'a, T> {
 188     #[inline]
 189     fn next(&mut self) -> Option<(&'a T)> { self.iter.next() }
 190
 191     #[inline]
 192     fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
 193 }
 194
 195 impl<T: Ord> FromIterator<T> for PriorityQueue<T> {
 196     fn from_iterator<Iter: Iterator<T>>(iter: Iter) -> PriorityQueue<T> {
 197         let mut q = PriorityQueue::new();
 198         q.extend(iter);
 199         q
 200     }
 201 }
 202
 203 impl<T: Ord> Extendable<T> for PriorityQueue<T> {
 204     fn extend<Iter: Iterator<T>>(&mut self, mut iter: Iter) {
 205         let (lower, _) = iter.size_hint();
 206
 207         let len = self.capacity();
 208         self.reserve(len + lower);
 209
 210         for elem in iter {
 211             self.push(elem);
 212         }
 213     }
 214 }
 215
 216 #[cfg(test)]
 217 mod tests {
 218     use priority_queue::PriorityQueue;
 219
 220     #[test]
 221     fn test_iterator() {
 222         let data = ~[5, 9, 3];
 223         let iterout = ~[9, 5, 3];
 224         let pq = PriorityQueue::from_vec(data);
 225         let mut i = 0;
 226         for el in pq.iter() {
 227             assert_eq!(*el, iterout[i]);
 228             i += 1;
 229         }
 230     }
 231
 232     #[test]
 233     fn test_top_and_pop() {
 234         let data = ~[2u, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1];
 235         let mut sorted = data.clone();
 236         sorted.sort();
 237         let mut heap = PriorityQueue::from_vec(data);
 238         while !heap.is_empty() {
 239             assert_eq!(heap.top(), sorted.last().unwrap());
 240             assert_eq!(heap.pop(), sorted.pop().unwrap());
 241         }
 242     }
 243
 244     #[test]
 245     fn test_push() {
 246         let mut heap = PriorityQueue::from_vec(~[2, 4, 9]);
 247         assert_eq!(heap.len(), 3);
 248         assert!(*heap.top() == 9);
 249         heap.push(11);
 250         assert_eq!(heap.len(), 4);
 251         assert!(*heap.top() == 11);
 252         heap.push(5);
 253         assert_eq!(heap.len(), 5);
 254         assert!(*heap.top() == 11);
 255         heap.push(27);
 256         assert_eq!(heap.len(), 6);
 257         assert!(*heap.top() == 27);
 258         heap.push(3);
 259         assert_eq!(heap.len(), 7);
 260         assert!(*heap.top() == 27);
 261         heap.push(103);
 262         assert_eq!(heap.len(), 8);
 263         assert!(*heap.top() == 103);
 264     }
 265
 266     #[test]
 267     fn test_push_unique() {
 268         let mut heap = PriorityQueue::from_vec(~[~2, ~4, ~9]);
 269         assert_eq!(heap.len(), 3);
 270         assert!(*heap.top() == ~9);
 271         heap.push(~11);
 272         assert_eq!(heap.len(), 4);
 273         assert!(*heap.top() == ~11);
 274         heap.push(~5);
 275         assert_eq!(heap.len(), 5);
 276         assert!(*heap.top() == ~11);
 277         heap.push(~27);
 278         assert_eq!(heap.len(), 6);
 279         assert!(*heap.top() == ~27);
 280         heap.push(~3);
 281         assert_eq!(heap.len(), 7);
 282         assert!(*heap.top() == ~27);
 283         heap.push(~103);
 284         assert_eq!(heap.len(), 8);
 285         assert!(*heap.top() == ~103);
 286     }
 287
 288     #[test]
 289     fn test_push_pop() {
 290         let mut heap = PriorityQueue::from_vec(~[5, 5, 2, 1, 3]);
 291         assert_eq!(heap.len(), 5);
 292         assert_eq!(heap.push_pop(6), 6);
 293         assert_eq!(heap.len(), 5);
 294         assert_eq!(heap.push_pop(0), 5);
 295         assert_eq!(heap.len(), 5);
 296         assert_eq!(heap.push_pop(4), 5);
 297         assert_eq!(heap.len(), 5);
 298         assert_eq!(heap.push_pop(1), 4);
 299         assert_eq!(heap.len(), 5);
 300     }
 301
 302     #[test]
 303     fn test_replace() {
 304         let mut heap = PriorityQueue::from_vec(~[5, 5, 2, 1, 3]);
 305         assert_eq!(heap.len(), 5);
 306         assert_eq!(heap.replace(6), 5);
 307         assert_eq!(heap.len(), 5);
 308         assert_eq!(heap.replace(0), 6);
 309         assert_eq!(heap.len(), 5);
 310         assert_eq!(heap.replace(4), 5);
 311         assert_eq!(heap.len(), 5);
 312         assert_eq!(heap.replace(1), 4);
 313         assert_eq!(heap.len(), 5);
 314     }
 315
 316     fn check_to_vec(mut data: ~[int]) {
 317         let heap = PriorityQueue::from_vec(data.clone());
 318         let mut v = heap.clone().to_vec();
 319         v.sort();
 320         data.sort();
 321
 322         assert_eq!(v, data);
 323         assert_eq!(heap.to_sorted_vec(), data);
 324     }
 325
 326     #[test]
 327     fn test_to_vec() {
 328         check_to_vec(~[]);
 329         check_to_vec(~[5]);
 330         check_to_vec(~[3, 2]);
 331         check_to_vec(~[2, 3]);
 332         check_to_vec(~[5, 1, 2]);
 333         check_to_vec(~[1, 100, 2, 3]);
 334         check_to_vec(~[1, 3, 5, 7, 9, 2, 4, 6, 8, 0]);
 335         check_to_vec(~[2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]);
 336         check_to_vec(~[9, 11, 9, 9, 9, 9, 11, 2, 3, 4, 11, 9, 0, 0, 0, 0]);
 337         check_to_vec(~[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
 338         check_to_vec(~[10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);
 339         check_to_vec(~[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0, 1, 2]);
 340         check_to_vec(~[5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1]);
 341     }
 342
 343     #[test]
 344     #[should_fail]
 345     fn test_empty_pop() {
 346         let mut heap: PriorityQueue<int> = PriorityQueue::new();
 347         heap.pop();
 348     }
 349
 350     #[test]
 351     fn test_empty_maybe_pop() {
 352         let mut heap: PriorityQueue<int> = PriorityQueue::new();
 353         assert!(heap.maybe_pop().is_none());
 354     }
 355
 356     #[test]
 357     #[should_fail]
 358     fn test_empty_top() {
 359         let empty: PriorityQueue<int> = PriorityQueue::new();
 360         empty.top();
 361     }
 362
 363     #[test]
 364     fn test_empty_maybe_top() {
 365         let empty: PriorityQueue<int> = PriorityQueue::new();
 366         assert!(empty.maybe_top().is_none());
 367     }
 368
 369     #[test]
 370     #[should_fail]
 371     fn test_empty_replace() {
 372         let mut heap: PriorityQueue<int> = PriorityQueue::new();
 373         heap.replace(5);
 374     }
 375
 376     #[test]
 377     fn test_from_iter() {
 378         let xs = ~[9u, 8, 7, 6, 5, 4, 3, 2, 1];
 379
 380         let mut q: PriorityQueue<uint> = xs.rev_iter().map(|&x| x).collect();
 381
 382         for &x in xs.iter() {
 383             assert_eq!(q.pop(), x);
 384         }
 385     }
 386 }