Skip to content

Commit

Permalink
Improve Heap Implementation (#785)
Browse files Browse the repository at this point in the history 
* ref: add `from_vec` method to create heap from vector

* chore: update docstrings
  • Loading branch information
@sozelfist
sozelfist committed Sep 10, 2024
1 parent 77da9fe commit 84c4665
Showing 1 changed file with 199 additions and 65 deletions.
264 changes: 199 additions & 65 deletions src/data_structures/heap.rs
View file
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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
}
Expand All @@ -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();
Expand All @@ -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]
Expand All @@ -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]
Expand All @@ -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));
}
}

0 comments on commit 84c4665

Please sign in to comment.