aoc/util/
heap.rs

1//! [Min heap] more suitable for algorithms such as [Dijkstra] and [A*] than Rust's default
2//! max heap. Splits the sorting key and value, so that you can order items without having
3//! to implement [`Ord`] on the value type.
4//!
5//! [Min heap]: https://en.wikipedia.org/wiki/Heap_(data_structure)
6//! [Dijkstra]: https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm
7//! [A*]: https://en.wikipedia.org/wiki/A*_search_algorithm
8use std::cmp::Ordering;
9use std::collections::BinaryHeap;
10
11struct Wrapper<K: Ord, V> {
12    key: K,
13    value: V,
14}
15
16impl<K: Ord, V> PartialEq for Wrapper<K, V> {
17    #[inline]
18    fn eq(&self, other: &Self) -> bool {
19        self.key == other.key
20    }
21}
22
23impl<K: Ord, V> Eq for Wrapper<K, V> {}
24
25impl<K: Ord, V> PartialOrd for Wrapper<K, V> {
26    #[inline]
27    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
28        Some(self.cmp(other))
29    }
30}
31
32impl<K: Ord, V> Ord for Wrapper<K, V> {
33    #[inline]
34    fn cmp(&self, other: &Self) -> Ordering {
35        other.key.cmp(&self.key)
36    }
37}
38
39#[derive(Default)]
40pub struct MinHeap<K: Ord, V> {
41    heap: BinaryHeap<Wrapper<K, V>>,
42}
43
44impl<K: Ord, V> MinHeap<K, V> {
45    #[must_use]
46    pub fn new() -> Self {
47        MinHeap { heap: BinaryHeap::new() }
48    }
49
50    #[must_use]
51    pub fn with_capacity(capacity: usize) -> Self {
52        MinHeap { heap: BinaryHeap::with_capacity(capacity) }
53    }
54
55    #[inline]
56    pub fn push(&mut self, key: K, value: V) {
57        self.heap.push(Wrapper { key, value });
58    }
59
60    #[inline]
61    #[must_use]
62    pub fn pop(&mut self) -> Option<(K, V)> {
63        self.heap.pop().map(|w| (w.key, w.value))
64    }
65}