algorithms_edu/data_structures/priority_queue/
binary_heap.rs

1// https://www.youtube.com/watch?v=QOJ-CmQiXko&list=PLDV1Zeh2NRsB6SWUrDFW2RmDotAfPbeHu&index=16
2
3use super::PriorityQueue;
4
5pub struct BinaryHeap<T: PartialOrd> {
6    heap: Vec<T>,
7}
8impl<T: PartialOrd> PriorityQueue<T> for BinaryHeap<T> {
9    fn with_capacity(sz: usize) -> Self {
10        Self {
11            heap: Vec::with_capacity(sz),
12        }
13    }
14    /// Adds an element to the priority queue, O(log(n))
15    fn insert(&mut self, el: T) {
16        self.heap.push(el);
17        self.swim(self.heap.len() - 1);
18    }
19    /// Test if an element is in heap, O(n)
20    fn contains(&self, el: &T) -> bool {
21        self.heap.contains(el)
22    }
23    /// Removes a particular element in the heap, O(n)
24    fn remove(&mut self, el: &T) {
25        if let Some(idx) = self.heap.iter().position(|x| x == el) {
26            self.remove_at(idx);
27        }
28    }
29    /// Removes the root of the heap, O(log(n))
30    fn poll(&mut self) -> Option<T> {
31        self.remove_at(0)
32    }
33}
34
35impl<T: PartialOrd> BinaryHeap<T> {
36    pub fn swap(&mut self, i: usize, j: usize) {
37        self.heap.swap(i, j);
38    }
39
40    /// Removes a node at particular index, O(log(n))
41    fn remove_at(&mut self, i: usize) -> Option<T> {
42        let end = self.heap.len() - 1;
43        self.heap.swap(i, end);
44        let removed = self.heap.pop();
45        // Try sinking element
46        let i_ = self.sink(i);
47        // If sinking did not work try swimming
48        if i_ == i {
49            self.swim(i);
50        }
51        removed
52    }
53
54    /// Perform bottom up node swim, O(log(n))
55    fn swim(&mut self, mut k: usize) -> usize {
56        // Grab the index of the next parent node WRT to k
57        let mut parent = (k.saturating_sub(1)) / 2;
58
59        // Keep swimming while we have not reached the
60        // root and while we're less than our parent.
61        while k > 0 && self.heap[k] < self.heap[parent] {
62            // Exchange k with the parent
63            self.heap.swap(parent, k);
64            k = parent;
65
66            // Grab the index of the next parent node WRT to k
67            parent = (k.saturating_sub(1)) / 2;
68        }
69        k
70    }
71
72    // Top down node sink, O(log(n))
73    fn sink(&mut self, mut k: usize) -> usize {
74        let heap_size = self.heap.len();
75        loop {
76            let left = 2 * k + 1; // Left  node
77            let right = 2 * k + 2; // Right node
78
79            // Find which is smaller left or right
80            let smallest = if right < heap_size && self.heap[right] < self.heap[left] {
81                right
82            } else {
83                left
84            };
85
86            // Stop if we're outside the bounds of the tree
87            // or stop early if we cannot sink k anymore
88            if left >= heap_size || self.heap[k] < self.heap[smallest] {
89                break;
90            }
91
92            // Move down the tree following the smallest node
93            self.heap.swap(smallest, k);
94            k = smallest;
95        }
96        k
97    }
98}
99#[cfg(test)]
100mod tests {
101    use super::*;
102    #[test]
103    fn test_priority_queue_binary_heap() {
104        let mut pq = BinaryHeap::with_capacity(8);
105        pq.insert(5);
106        pq.insert(7);
107        pq.insert(3);
108        pq.insert(8);
109        pq.insert(2);
110        pq.insert(1);
111        assert_eq!(pq.poll().unwrap(), 1);
112        pq.remove(&2);
113        assert_eq!(pq.poll().unwrap(), 3);
114    }
115}
algorithms_edu/data_structures/priority_queue/binary_heap.rs

algorithms_edu/data_structures/priority_queue/
binary_heap.rs
