MinHeap utility class

Author: maneatingape

src/lib.rs

@@ -148,6 +148,7 @@ pub mod util {
     pub mod ansi;
     pub mod grid;
     pub mod hash;
+    pub mod heap;
     pub mod integer;
     pub mod iter;
     pub mod math;

src/util/heap.rs

@@ -0,0 +1,62 @@
+//! [Min heap] more suitable for algorithms such as [Dijkstra] and [A*] than Rust's default
+//! max heap. Splits the sorting key and value, so that you can order items without having
+//! to implement [`Ord`] on the value type.
+//!
+//! [Min heap]: https://en.wikipedia.org/wiki/Heap_(data_structure)
+//! [Dijkstra]: https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm
+//! [A*]: https://en.wikipedia.org/wiki/A*_search_algorithm
+use std::cmp::Ordering;
+use std::collections::BinaryHeap;
+
+struct Wrapper<K: Ord, V> {
+    key: K,
+    value: V,
+}
+
+impl<K: Ord, V> PartialEq for Wrapper<K, V> {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.key == other.key
+    }
+}
+
+impl<K: Ord, V> Eq for Wrapper<K, V> {}
+
+impl<K: Ord, V> PartialOrd for Wrapper<K, V> {
+    #[inline]
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<K: Ord, V> Ord for Wrapper<K, V> {
+    #[inline]
+    fn cmp(&self, other: &Self) -> Ordering {
+        other.key.cmp(&self.key)
+    }
+}
+
+#[derive(Default)]
+pub struct MinHeap<K: Ord, V> {
+    heap: BinaryHeap<Wrapper<K, V>>,
+}
+
+impl<K: Ord, V> MinHeap<K, V> {
+    pub fn new() -> Self {
+        MinHeap { heap: BinaryHeap::new() }
+    }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        MinHeap { heap: BinaryHeap::with_capacity(capacity) }
+    }
+
+    #[inline]
+    pub fn push(&mut self, key: K, value: V) {
+        self.heap.push(Wrapper { key, value });
+    }
+
+    #[inline]
+    pub fn pop(&mut self) -> Option<(K, V)> {
+        self.heap.pop().map(|w| (w.key, w.value))
+    }
+}

src/year2021/day23.rs

@@ -1,7 +1,6 @@
 use crate::util::hash::*;
+use crate::util::heap::*;
 use std::array::from_fn;
-use std::cmp::Ordering;
-use std::collections::BinaryHeap;
 use std::hash::*;
 
 const A: usize = 0;
@@ -88,26 +87,6 @@ impl Burrow {
     }
 }
 
-#[derive(Copy, Clone, PartialEq, Eq)]
-struct State {
-    burrow: Burrow,
-    energy: usize,
-}
-
-impl PartialOrd for State {
-    #[inline]
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for State {
-    #[inline]
-    fn cmp(&self, other: &Self) -> Ordering {
-        other.energy.cmp(&self.energy)
-    }
-}
-
 pub fn parse(input: &str) -> Vec<Vec<usize>> {
     input
         .lines()
@@ -136,13 +115,12 @@ pub fn part2(input: &[Vec<usize>]) -> usize {
 }
 
 fn organize(burrow: Burrow) -> usize {
-    let mut todo = BinaryHeap::with_capacity(20_000);
+    let mut todo = MinHeap::with_capacity(20_000);
     let mut seen = FastMap::with_capacity(20_000);
 
-    todo.push(State { burrow, energy: best_possible(&burrow) });
+    todo.push(best_possible(&burrow), burrow);
 
-    while let Some(state) = todo.pop() {
-        let State { mut burrow, energy } = state;
+    while let Some((energy, mut burrow)) = todo.pop() {
         let open: [bool; 4] = from_fn(|i| burrow.rooms[i].open(i));
 
         let mut changed = false;
@@ -164,7 +142,7 @@ fn organize(burrow: Burrow) -> usize {
             // Moving back to home burrow does not change cost
             let min = seen.get(&burrow).unwrap_or(&usize::MAX);
             if energy < *min {
-                todo.push(State { burrow, energy });
+                todo.push(energy, burrow);
                 seen.insert(burrow, energy);
             }
         } else {
@@ -173,8 +151,8 @@ fn organize(burrow: Burrow) -> usize {
                     let offset = 2 + 2 * i;
                     let forward = (offset + 1)..11;
                     let reverse = (0..offset).rev();
-                    expand(&mut todo, &mut seen, &state, i, forward);
-                    expand(&mut todo, &mut seen, &state, i, reverse);
+                    expand(&mut todo, &mut seen, burrow, energy, i, forward);
+                    expand(&mut todo, &mut seen, burrow, energy, i, reverse);
                 }
             }
         }
@@ -240,13 +218,13 @@ fn condense(burrow: &mut Burrow, kind: usize, iter: impl Iterator<Item = usize>)
 }
 
 fn expand(
-    todo: &mut BinaryHeap<State>,
+    todo: &mut MinHeap<usize, Burrow>,
     seen: &mut FastMap<Burrow, usize>,
-    state: &State,
+    mut burrow: Burrow,
+    energy: usize,
     room_index: usize,
     iter: impl Iterator<Item = usize>,
 ) {
-    let State { mut burrow, energy } = state;
     let kind = burrow.rooms[room_index].pop();
 
     for hallway_index in iter {
@@ -290,7 +268,7 @@ fn expand(
                 let min = seen.get(&next).unwrap_or(&usize::MAX);
 
                 if next_energy < *min {
-                    todo.push(State { burrow: next, energy: next_energy });
+                    todo.push(next_energy, next);
                     seen.insert(next, next_energy);
                 }
             }