diff --git a/src/data_structures/heap.rs b/src/data_structures/heap.rs index 2d882f7af48..cb48ff1bbd1 100644 --- a/src/data_structures/heap.rs +++ b/src/data_structures/heap.rs @@ -1,86 +1,160 @@ -// Heap data structure -// Takes a closure as a comparator to allow for min-heap, max-heap, and works with custom key functions +//! A generic heap data structure. +//! +//! This module provides a `Heap` implementation that can function as either a +//! min-heap or a max-heap. It supports common heap operations such as adding, +//! removing, and iterating over elements. The heap can also be created from +//! an unsorted vector and supports custom comparators for flexible sorting +//! behavior. use std::{cmp::Ord, slice::Iter}; +/// A heap data structure that can be used as a min-heap, max-heap or with +/// custom comparators. +/// +/// This struct manages a collection of items where the heap property is maintained. +/// The heap can be configured to order elements based on a provided comparator function, +/// allowing for both min-heap and max-heap functionalities, as well as custom sorting orders. pub struct Heap { items: Vec, comparator: fn(&T, &T) -> bool, } impl Heap { + /// Creates a new, empty heap with a custom comparator function. + /// + /// # Parameters + /// - `comparator`: A function that defines the heap's ordering. + /// + /// # Returns + /// A new `Heap` instance. pub fn new(comparator: fn(&T, &T) -> bool) -> Self { Self { - // Add a default in the first spot to offset indexes - // for the parent/child math to work out. - // Vecs have to have all the same type so using Default - // is a way to add an unused item. items: vec![], comparator, } } + /// Creates a heap from a vector and a custom comparator function. + /// + /// # Parameters + /// - `items`: A vector of items to be turned into a heap. + /// - `comparator`: A function that defines the heap's ordering. + /// + /// # Returns + /// A `Heap` instance with the elements from the provided vector. + pub fn from_vec(items: Vec, comparator: fn(&T, &T) -> bool) -> Self { + let mut heap = Self { items, comparator }; + heap.build_heap(); + heap + } + + /// Constructs the heap from an unsorted vector by applying the heapify process. + fn build_heap(&mut self) { + let last_parent_idx = (self.len() / 2).wrapping_sub(1); + for idx in (0..=last_parent_idx).rev() { + self.heapify_down(idx); + } + } + + /// Returns the number of elements in the heap. + /// + /// # Returns + /// The number of elements in the heap. pub fn len(&self) -> usize { self.items.len() } + /// Checks if the heap is empty. + /// + /// # Returns + /// `true` if the heap is empty, `false` otherwise. pub fn is_empty(&self) -> bool { self.len() == 0 } + /// Adds a new element to the heap and maintains the heap property. + /// + /// # Parameters + /// - `value`: The value to add to the heap. pub fn add(&mut self, value: T) { self.items.push(value); - - // Heapify Up - let mut idx = self.len() - 1; - while let Some(pdx) = self.parent_idx(idx) { - if (self.comparator)(&self.items[idx], &self.items[pdx]) { - self.items.swap(idx, pdx); - } - idx = pdx; - } + self.heapify_up(self.len() - 1); } + /// Removes and returns the root element from the heap. + /// + /// # Returns + /// The root element if the heap is not empty, otherwise `None`. pub fn pop(&mut self) -> Option { if self.is_empty() { return None; } - // This feels like a function built for heap impl :) - // Removes an item at an index and fills in with the last item - // of the Vec let next = Some(self.items.swap_remove(0)); - if !self.is_empty() { - // Heapify Down - let mut idx = 0; - while self.children_present(idx) { - let cdx = { - if self.right_child_idx(idx) >= self.len() { - self.left_child_idx(idx) - } else { - let ldx = self.left_child_idx(idx); - let rdx = self.right_child_idx(idx); - if (self.comparator)(&self.items[ldx], &self.items[rdx]) { - ldx - } else { - rdx - } - } - }; - if !(self.comparator)(&self.items[idx], &self.items[cdx]) { - self.items.swap(idx, cdx); - } - idx = cdx; - } + self.heapify_down(0); } - next } + /// Returns an iterator over the elements in the heap. + /// + /// # Returns + /// An iterator over the elements in the heap, in their internal order. pub fn iter(&self) -> Iter<'_, T> { self.items.iter() } + /// Moves an element upwards to restore the heap property. + /// + /// # Parameters + /// - `idx`: The index of the element to heapify up. + fn heapify_up(&mut self, mut idx: usize) { + while let Some(pdx) = self.parent_idx(idx) { + if (self.comparator)(&self.items[idx], &self.items[pdx]) { + self.items.swap(idx, pdx); + idx = pdx; + } else { + break; + } + } + } + + /// Moves an element downwards to restore the heap property. + /// + /// # Parameters + /// - `idx`: The index of the element to heapify down. + fn heapify_down(&mut self, mut idx: usize) { + while self.children_present(idx) { + let cdx = { + if self.right_child_idx(idx) >= self.len() { + self.left_child_idx(idx) + } else { + let ldx = self.left_child_idx(idx); + let rdx = self.right_child_idx(idx); + if (self.comparator)(&self.items[ldx], &self.items[rdx]) { + ldx + } else { + rdx + } + } + }; + + if (self.comparator)(&self.items[cdx], &self.items[idx]) { + self.items.swap(idx, cdx); + idx = cdx; + } else { + break; + } + } + } + + /// Returns the index of the parent of the element at `idx`. + /// + /// # Parameters + /// - `idx`: The index of the element. + /// + /// # Returns + /// The index of the parent element if it exists, otherwise `None`. fn parent_idx(&self, idx: usize) -> Option { if idx > 0 { Some((idx - 1) / 2) @@ -89,14 +163,35 @@ impl Heap { } } + /// Checks if the element at `idx` has children. + /// + /// # Parameters + /// - `idx`: The index of the element. + /// + /// # Returns + /// `true` if the element has children, `false` otherwise. fn children_present(&self, idx: usize) -> bool { - self.left_child_idx(idx) <= (self.len() - 1) + self.left_child_idx(idx) < self.len() } + /// Returns the index of the left child of the element at `idx`. + /// + /// # Parameters + /// - `idx`: The index of the element. + /// + /// # Returns + /// The index of the left child. fn left_child_idx(&self, idx: usize) -> usize { idx * 2 + 1 } + /// Returns the index of the right child of the element at `idx`. + /// + /// # Parameters + /// - `idx`: The index of the element. + /// + /// # Returns + /// The index of the right child. fn right_child_idx(&self, idx: usize) -> usize { self.left_child_idx(idx) + 1 } @@ -106,20 +201,49 @@ impl Heap where T: Ord, { - /// Create a new MinHeap + /// Creates a new min-heap. + /// + /// # Returns + /// A new `Heap` instance configured as a min-heap. pub fn new_min() -> Heap { Self::new(|a, b| a < b) } - /// Create a new MaxHeap + /// Creates a new max-heap. + /// + /// # Returns + /// A new `Heap` instance configured as a max-heap. pub fn new_max() -> Heap { Self::new(|a, b| a > b) } + + /// Creates a min-heap from an unsorted vector. + /// + /// # Parameters + /// - `items`: A vector of items to be turned into a min-heap. + /// + /// # Returns + /// A `Heap` instance configured as a min-heap. + pub fn from_vec_min(items: Vec) -> Heap { + Self::from_vec(items, |a, b| a < b) + } + + /// Creates a max-heap from an unsorted vector. + /// + /// # Parameters + /// - `items`: A vector of items to be turned into a max-heap. + /// + /// # Returns + /// A `Heap` instance configured as a max-heap. + pub fn from_vec_max(items: Vec) -> Heap { + Self::from_vec(items, |a, b| a > b) + } } #[cfg(test)] mod tests { use super::*; + #[test] fn test_empty_heap() { let mut heap: Heap = Heap::new_max(); @@ -139,6 +263,8 @@ mod tests { assert_eq!(heap.pop(), Some(9)); heap.add(1); assert_eq!(heap.pop(), Some(1)); + assert_eq!(heap.pop(), Some(11)); + assert_eq!(heap.pop(), None); } #[test] @@ -154,22 +280,8 @@ mod tests { assert_eq!(heap.pop(), Some(4)); heap.add(1); assert_eq!(heap.pop(), Some(2)); - } - - #[allow(dead_code)] - struct Point(/* x */ i32, /* y */ i32); - - #[test] - fn test_key_heap() { - let mut heap: Heap = Heap::new(|a, b| a.0 < b.0); - heap.add(Point(1, 5)); - heap.add(Point(3, 10)); - heap.add(Point(-2, 4)); - assert_eq!(heap.len(), 3); - assert_eq!(heap.pop().unwrap().0, -2); - assert_eq!(heap.pop().unwrap().0, 1); - heap.add(Point(50, 34)); - assert_eq!(heap.pop().unwrap().0, 3); + assert_eq!(heap.pop(), Some(1)); + assert_eq!(heap.pop(), None); } #[test] @@ -180,15 +292,13 @@ mod tests { heap.add(9); heap.add(11); - // test iterator, which is not in order except the first one. let mut iter = heap.iter(); assert_eq!(iter.next(), Some(&2)); - assert_ne!(iter.next(), None); - assert_ne!(iter.next(), None); - assert_ne!(iter.next(), None); + assert_eq!(iter.next(), Some(&4)); + assert_eq!(iter.next(), Some(&9)); + assert_eq!(iter.next(), Some(&11)); assert_eq!(iter.next(), None); - // test the heap after run iterator. assert_eq!(heap.len(), 4); assert_eq!(heap.pop(), Some(2)); assert_eq!(heap.pop(), Some(4)); @@ -196,4 +306,28 @@ mod tests { assert_eq!(heap.pop(), Some(11)); assert_eq!(heap.pop(), None); } + + #[test] + fn test_from_vec_min() { + let vec = vec![3, 1, 4, 1, 5, 9, 2, 6, 5]; + let mut heap = Heap::from_vec_min(vec); + assert_eq!(heap.len(), 9); + assert_eq!(heap.pop(), Some(1)); + assert_eq!(heap.pop(), Some(1)); + assert_eq!(heap.pop(), Some(2)); + heap.add(0); + assert_eq!(heap.pop(), Some(0)); + } + + #[test] + fn test_from_vec_max() { + let vec = vec![3, 1, 4, 1, 5, 9, 2, 6, 5]; + let mut heap = Heap::from_vec_max(vec); + assert_eq!(heap.len(), 9); + assert_eq!(heap.pop(), Some(9)); + assert_eq!(heap.pop(), Some(6)); + assert_eq!(heap.pop(), Some(5)); + heap.add(10); + assert_eq!(heap.pop(), Some(10)); + } }