Improve Heap Implementation (#785)

* ref: add `from_vec` method to create heap from vector

* chore: update docstrings

Author: sozelfist

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<T> {
     items: Vec<T>,
     comparator: fn(&T, &T) -> bool,
 }
 
 impl<T> Heap<T> {
+    /// 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<T>, 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<T> {
         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<usize> {
         if idx > 0 {
             Some((idx - 1) / 2)
@@ -89,14 +163,35 @@ impl<T> Heap<T> {
         }
     }
 
+    /// 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<T> Heap<T>
 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<T> {
         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<T> {
         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<T>) -> Heap<T> {
+        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<T>) -> Heap<T> {
+        Self::from_vec(items, |a, b| a > b)
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+
     #[test]
     fn test_empty_heap() {
         let mut heap: Heap<i32> = 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<Point> = 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,20 +292,42 @@ 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));
         assert_eq!(heap.pop(), Some(9));
         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));
+    }
 }