tree-experiments 0.1.0

// SPDX-FileCopyrightText: 2023 Markus Mayer
// SPDX-License-Identifier: EUPL-1.2

use std::fmt::{Debug, Formatter};
use std::marker::PhantomData;
use std::mem::swap;

/// A binary max-heap.
///
/// Uses a list representation internally.
pub type BinaryMaxHeap<T> = BinaryHeap<T, Max<T>>;

/// A binary max-heap.
///
/// Uses a list representation internally.
pub type BinaryMinHeap<T> = BinaryHeap<T, Min<T>>;

/// A binary tree with an arbitrary heap property.
pub struct BinaryHeap<T, H> {
    /// The heap in list representation.
    ///
    /// ## List representation
    ///
    /// An element at index `N` has its left child at index `2×N+1` and its right child at `2×N+2`.
    /// The parent of a non-root element at `N>0` can be found at index `(N-1)/2`.
    ///
    /// ## Complete Binary Tree
    ///
    /// The heap is implemented as a complete binary tree, i.e. all levels are completely filled,
    /// except for possibly the last level, i.e. all children are as far left as possible, and a
    /// node has either both a left and right child, just a left child or no children at all.
    data: Vec<T>,
    _heap_property: PhantomData<H>,
}

/// Implements the max-heap property.
pub struct Max<T>(PhantomData<T>);

/// Implements the min-heap property.
pub struct Min<T>(PhantomData<T>);

/// Heap property.
pub trait HeapProperty<T> {
    const NAME: &'static str;

    /// Determines whether the heap property is upheld between the first and the
    /// second value, where the first value takes priority:
    ///
    /// * `property_upheld(parent, child)`
    /// * `property_upheld(left, right)`
    fn property_upheld(lhs: &T, rhs: &T) -> bool;
}

impl<T, H> BinaryHeap<T, H> {
    /// Constructs a new, empty `BinaryHeap`.
    pub const fn new() -> Self {
        Self {
            data: Vec::new(),
            _heap_property: PhantomData,
        }
    }

    /// Constructs a new, empty `BinaryHeap` with at least the specified capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            data: Vec::with_capacity(capacity),
            _heap_property: PhantomData,
        }
    }

    /// Pushes an element onto the heap.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to push onto the heap.
    pub fn push(&mut self, value: T)
    where
        T: PartialOrd,
        H: HeapProperty<T>,
    {
        let index = self.data.len();
        self.data.push(value);
        self.upheap(index);
    }

    /// Gets a reference to the top element on the heap.
    ///
    /// ## Returns
    ///
    /// A value if the tree is nonempty; or `None` otherwise.
    pub fn peek(&self) -> Option<&T> {
        self.data.get(0)
    }

    /// Extracts an element from the heap.
    ///
    /// ## Returns
    ///
    /// A value if the tree is nonempty; or `None` otherwise.
    pub fn pop(&mut self) -> Option<T>
    where
        T: PartialOrd,
        H: HeapProperty<T>,
    {
        if self.data.is_empty() {
            return None;
        }

        // Replace the root of the heap with the last element on the last level.
        let last_index = self.data.len() - 1;
        self.data.swap(0, last_index);

        // We extract the last element first so it doesn't interfere with the downheap operation.
        let extracted_value = self.data.pop();

        if !self.data.is_empty() {
            self.downheap(0);
        }

        extracted_value
    }

    /// Pushes an element onto the heap and extracts the top value in one operation.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to push onto the heap.
    ///
    /// ## Returns
    ///
    /// The value. If the tree is empty, the value itself is returned.
    pub fn push_pop(&mut self, mut value: T) -> T
    where
        T: PartialOrd,
        H: HeapProperty<T>,
    {
        if self.data.is_empty() {
            return value;
        }

        if H::property_upheld(&value, &self.data[0]) {
            return value;
        }

        swap(&mut self.data[0], &mut value);
        self.downheap(0);

        value
    }

    /// Determines whether the heap contains the specified item.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to test for.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the heap contains the specified item or `false` if not.
    pub fn contains(&self, value: &T) -> bool
    where
        T: PartialEq,
    {
        self.data.contains(value)
    }

    /// Deletes a value from the heap.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to delete.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the value was deleted or `false` if not.
    pub fn remove(&mut self, value: &T) -> bool
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
    {
        self.remove_where(|x| x == value).is_some()
    }

    /// Deletes a value from the heap.
    ///
    /// ## Arguments
    ///
    /// * `predicate` - The predicate used to test for the value. If the predicate matches,
    ///                 the value is removed and the function returns.
    ///
    /// ## Returns
    ///
    /// Returns `Some(value)` if the value was deleted or `None` if not.
    pub fn remove_where<F>(&mut self, predicate: F) -> Option<T>
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
        F: FnMut(&T) -> bool,
    {
        if self.data.is_empty() {
            return None;
        }

        let last_index = self.data.len() - 1;

        // Find the index of the element we want to delete.
        if let Some(index) = self.data.iter().position(predicate) {
            // Swap this element with the last element.
            self.data.swap(index, last_index);

            // Remove the element.
            let previous_value = self.data.pop().expect("the value exists");

            // If we deleted the last element, exit.
            if index == last_index {
                return Some(previous_value);
            }

            // Down-heapify or up-heapify to restore the heap property.
            let up_heapify = H::property_upheld(&self.data[index], &previous_value);
            if up_heapify {
                self.upheap(index);
            } else {
                self.downheap(index);
            }

            Some(previous_value)
        } else {
            None
        }
    }

    /// Replaces a value on the heap.
    ///
    /// This function can be used if the replacement value already exists. If the element needs
    /// to be mutated in place, use [`BinaryHeap::replace_where_with`] instead.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to replace.
    /// * `replacement` - The replacement value.
    ///
    /// ## Returns
    ///
    /// Returns `Some(previous value)` if the value was replaced or `None` if not.
    pub fn replace(&mut self, value: &T, replacement: T) -> Option<T>
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
    {
        self.replace_where(|x| x == value, replacement)
    }

    /// Replaces a value on the heap.
    ///
    /// This function can be used if the replacement value already exists. If the element needs
    /// to be mutated in place, use [`BinaryHeap::replace_where_with`] instead.
    ///
    /// ## Arguments
    ///
    /// * `predicate` - The predicate used to test for the value. If the predicate matches,
    ///                 the value is removed and the function returns.
    /// * `replacement` - The replacement value.
    ///
    /// ## Returns
    ///
    /// Returns `Some(previous value)` if the value was replaced or `None` if not.
    pub fn replace_where<F>(&mut self, predicate: F, mut replacement: T) -> Option<T>
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
        F: FnMut(&T) -> bool,
    {
        if self.data.is_empty() {
            return None;
        }

        // Find the index of the element we want to delete.
        if let Some(index) = self.data.iter().position(predicate) {
            let downheap = H::property_upheld(&self.data[index], &replacement);

            swap(&mut self.data[index], &mut replacement);

            if downheap {
                self.downheap(index);
            } else {
                self.upheap(index);
            }

            Some(replacement)
        } else {
            None
        }
    }

    /// Replaces a value on the heap.
    ///
    /// This function allows to mutate the value in-place but at the cost of a bit of performance.
    /// If the replacement value exists beforehand, use [`BinaryHeap::replace_where`] instead.
    ///
    /// ## Arguments
    ///
    /// * `value` - The value to replace.
    /// * `replacement` - The function creating the replacement value.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the value was replaced or `false` if not.
    pub fn replace_with<R>(&mut self, value: &T, replacement: R) -> bool
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
        R: FnMut(&mut T),
    {
        self.replace_where_with(|x| x == value, replacement)
    }

    /// Replaces a value on the heap.
    ///
    /// This function allows to mutate the value in-place but at the cost of a bit of performance.
    /// If the replacement value exists beforehand, use [`BinaryHeap::replace_where`] instead.
    ///
    /// ## Arguments
    ///
    /// * `predicate` - The predicate used to test for the value. If the predicate matches,
    ///                 the value is removed and the function returns.
    /// * `replacement` - The function creating the replacement value.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the value was replaced or `false` if not.
    pub fn replace_where_with<F, R>(&mut self, predicate: F, mut replacement: R) -> bool
    where
        T: PartialEq + PartialOrd,
        H: HeapProperty<T>,
        F: FnMut(&T) -> bool,
        R: FnMut(&mut T),
    {
        if self.data.is_empty() {
            return false;
        }

        // Find the index of the element we want to delete.
        if let Some(index) = self.data.iter().position(predicate) {
            // Mutate the element.
            replacement(&mut self.data[index]);

            // TODO: Optimize these calls? We don't know if the value increased or decreased here.
            let _ = self.upheap(index) || self.downheap(index);

            true
        } else {
            false
        }
    }

    /// Returns the number of elements on the heap.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Returns `true` if the heap is empty.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Ensures that elements are in the correct order, starting from a given child node,
    /// working its way up the heap.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the tree was changed or `false` if no change was performed.
    fn upheap(&mut self, mut index: usize) -> bool
    where
        T: PartialOrd,
        H: HeapProperty<T>,
    {
        debug_assert!(index < self.data.len());
        let mut changed = false;

        // Compare the added element with its parent; if they are in the correct order, stop.
        while index > 0 {
            let parent_idx = self
                .get_parent_idx(index)
                .expect("since the tree was not empty before the insert, the parent must exist");

            let parent = self
                .data
                .get(parent_idx)
                .expect("since the tree was not empty before the insert, the parent must exist");

            if H::property_upheld(&parent, &self.data[index]) {
                break;
            }

            // If not, swap the element with its parent and return to the previous step.
            self.data.swap(parent_idx, index);
            index = parent_idx;

            changed = true;
        }

        changed
    }

    /// Ensures that elements are in the correct order, starting from a given parent node,
    /// working its way down the heap.
    ///
    /// ## Returns
    ///
    /// Returns `true` if the tree was changed or `false` if no change was performed.
    fn downheap(&mut self, mut index: usize) -> bool
    where
        T: PartialOrd,
        H: HeapProperty<T>,
    {
        debug_assert!(index < self.data.len());
        let mut changed = false;

        // Compare the new root with its children; if they are in the correct order, stop.
        // If not, swap the element with one of its children and return to the previous step.
        // (Swap with its smaller child in a min-heap and its larger child in a max-heap.)
        loop {
            let left_child_idx = self.get_left_child_idx(index);
            let right_child_idx = self.get_right_child_idx(index);

            let current = &self.data[index];

            match (left_child_idx, right_child_idx) {
                (Some(left), Some(right)) => {
                    let left_child = &self.data[left];
                    let right_child = &self.data[right];

                    if H::property_upheld(&current, left_child)
                        && H::property_upheld(current, right_child)
                    {
                        break;
                    }

                    if H::property_upheld(left_child, right_child) {
                        self.data.swap(index, left);
                        index = left;
                    } else {
                        self.data.swap(index, right);
                        index = right;
                    }

                    changed = true;
                }
                (Some(child), None) => {
                    let left_child = &self.data[child];

                    if H::property_upheld(current, left_child) {
                        break;
                    }

                    self.data.swap(index, child);

                    // Since there is no right child and the tree is balanced, the left child can also not have any children.
                    debug_assert!(self.get_left_child_idx(child).is_none());

                    changed = true;
                }
                (None, Some(_)) => unreachable!("the tree cannot have only a right child"),
                (None, None) => break,
            }
        }

        changed
    }

    /// Gets the index of the left child of the item at the specified `index`.
    const fn get_left_child_idx_unchecked(index: usize) -> usize {
        2 * index + 1
    }

    /// Gets the index of the left child of the item at the specified `index`.
    ///
    /// ## Returns
    ///
    /// * `Some(index)` if the element exists in the tree.
    /// * `None` if the child does not exist.
    fn get_left_child_idx(&self, index: usize) -> Option<usize> {
        let idx = Self::get_left_child_idx_unchecked(index);
        if idx < self.data.len() {
            Some(idx)
        } else {
            None
        }
    }

    /// Gets the index of the right child of the item at the specified `index`.
    const fn get_right_child_idx_unchecked(index: usize) -> usize {
        2 * index + 2
    }

    /// Gets the index of the right child of the item at the specified `index`.
    ///
    /// ## Returns
    ///
    /// * `Some(index)` if the element exists in the tree.
    /// * `None` if the child does not exist.
    fn get_right_child_idx(&self, index: usize) -> Option<usize> {
        let idx = Self::get_right_child_idx_unchecked(index);
        if idx < self.data.len() {
            Some(idx)
        } else {
            None
        }
    }

    /// Gets the index of the parent of the item at the specified `index`.
    ///
    /// ## Returns
    ///
    /// * `Some(index)` if the element exists in the tree.
    /// * `None` if the parent does not exist.
    fn get_parent_idx(&self, index: usize) -> Option<usize> {
        if index == 0 || index >= self.data.len() {
            return None;
        }

        Some((index - 1) / 2)
    }
}

impl<T, H> FromIterator<T> for BinaryHeap<T, H>
where
    T: PartialOrd,
    H: HeapProperty<T>,
{
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        <Self as From<Vec<T>>>::from(iter.into_iter().collect())
    }
}

impl<T, H> From<Vec<T>> for BinaryHeap<T, H>
where
    T: PartialOrd,
    H: HeapProperty<T>,
{
    fn from(value: Vec<T>) -> Self {
        let mut tree = Self {
            data: value,
            _heap_property: PhantomData,
        };

        // Floyd's method.
        let length = tree.data.len();
        for i in (0..=(length / 2)).rev() {
            tree.downheap(i);
        }

        tree
    }
}

impl<T, H> Debug for BinaryHeap<T, H>
where
    T: Debug,
    H: HeapProperty<T>,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct(&format!("Binary{}Heap", H::NAME))
            .field("data", &self.data)
            .finish()
    }
}

impl<T, H> Default for BinaryHeap<T, H> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T, H> Clone for BinaryHeap<T, H>
where
    T: Clone,
{
    fn clone(&self) -> Self {
        Self {
            data: self.data.clone(),
            _heap_property: PhantomData,
        }
    }
}

impl<T, H> AsRef<[T]> for BinaryHeap<T, H> {
    fn as_ref(&self) -> &[T] {
        &self.data
    }
}

impl<T> HeapProperty<T> for Max<T>
where
    T: PartialOrd,
{
    const NAME: &'static str = "Max";

    fn property_upheld(lhs: &T, rhs: &T) -> bool {
        lhs >= rhs
    }
}

impl<T> HeapProperty<T> for Min<T>
where
    T: PartialOrd,
{
    const NAME: &'static str = "Min";

    fn property_upheld(lhs: &T, rhs: &T) -> bool {
        lhs <= rhs
    }
}