pyority_queue/src/backing/pure/binary_heap.rs

use std::fmt;

use super::PureBacking;

/// A binary min-heap backed by an array
#[derive(Debug)]
pub struct BinaryHeap<T: Ord + Clone + Send + Sync> {
    data: Vec<T>,
}

/// Indicates why a sift failed
#[derive(Debug, Clone)]
struct SiftError {
    /// The index that couldn't be sifted
    index: usize,
    /// The length of the array
    len: usize,
}

impl SiftError {
    /// Instantiates a `SiftError`
    fn new(index: usize, len: usize) -> Self {
        Self { index, len }
    }
}

impl fmt::Display for SiftError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Could not sift index {} of {}", self.index, self.len)
    }
}

/// Whether a sift operation succeeded
type SiftResult = Result<(), SiftError>;

impl<T: Ord + Clone + Send + Sync> BinaryHeap<T> {
    /// Instantiates a new (empty) binary heap
    pub fn new() -> Self {
        Self { data: vec![] }
    }

    /// Fix an index representing a node with valid children but that may violate the heap property compared to its immediate parent
    fn sift_up(&mut self, i: usize) -> SiftResult {
        if i == 0 {
            // Base case, at root so nothing to do
            Ok(())
        } else if let Some(child) = self.data.get(i).cloned() {
            let parent_index = (i - 1) / 2;
            // Check if the heap property is violated
            if child < self.data[parent_index] {
                // Swap child with parent
                self.data[i] = self.data[parent_index].clone();
                self.data[parent_index] = child;

                // Repeat process with parent
                self.sift_up(parent_index)
            } else {
                // Sift complete
                Ok(())
            }
        } else {
            // Tried to sift a non-existent index
            Err(SiftError::new(i, self.data.len()))
        }
    }
    /// Fix an index representing a node with that doesn't violate the heap property, but may have immediate children that do
    fn sift_down(&mut self, i: usize) -> SiftResult {
        if i > self.data.len() {
            // Tried to sift a non-existent index
            Err(SiftError::new(i, self.data.len()))
        } else {
            if let Some(first_child) = self.data.get(i * 2 + 1).cloned() {
                let smaller_child_index;
                let smaller_child;

                // Find the smallest child and its index
                if let Some(second_child) = self.data.get(i * 2 + 2).cloned() {
                    // Both children, use the smaller one
                    if first_child < second_child {
                        smaller_child = first_child;
                        smaller_child_index = i * 2 + 1;
                    } else {
                        smaller_child = second_child;
                        smaller_child_index = i * 2 + 2;
                    }
                } else {
                    // Only one child, no choice
                    smaller_child = first_child;
                    smaller_child_index = i * 2 + 1;
                }

                if smaller_child < self.data[i] {
                    // Swap parent with child
                    self.data[smaller_child_index] = self.data[i].clone();
                    self.data[i] = smaller_child;

                    // Repeat process with child
                    self.sift_down(smaller_child_index)
                } else {
                    // Heap property satisfied, we're done
                    Ok(())
                }
            } else {
                // Base case, no children so nothing to do
                Ok(())
            }
        }
    }
}

impl<T: Ord + Clone + Send + Sync> FromIterator<T> for BinaryHeap<T> {
    fn from_iter<U: IntoIterator<Item = T>>(iter: U) -> Self {
        let mut this = Self {
            data: Vec::from_iter(iter),
        };
        for i in (0..=(this.data.len() / 2)).rev() {
            this.sift_down(i).expect("Index error during heapify");
        }
        this
    }
}

impl<T: Ord + Clone + Send + Sync> PureBacking<T> for BinaryHeap<T> {
    fn add(&mut self, item: T) {
        // Append item
        self.data.push(item);
        // Sift up to retain heap property
        self.sift_up(self.data.len() - 1).unwrap();
    }

    fn pop(&mut self) -> Option<T> {
        match self.data.len() {
            // No extra processing
            0 | 1 => self.data.pop(),
            _ => {
                // Get the original root
                let root = self.data[0].clone();

                // Move final item to the root and sift down to regain heap property
                self.data[0] = self
                    .data
                    .pop()
                    .expect("Vector claimed not to be empty but was");
                self.sift_down(0).unwrap();

                // Return original root
                Some(root)
            }
        }
    }

    fn len(&self) -> usize {
        self.data.len()
    }
}