pyority_queue/src/backing/indexed/binary_heap.rs

use std::{collections::HashMap, hash::Hash};

use crate::backing::{
    containers::{Pair, SiftError, SiftResult},
    indexed::IndexedBacking,
};

pub struct IndexedBinaryHeap<
    D: Hash + Eq + Clone + Send + Sync,
    P: PartialOrd + Clone + Send + Sync,
> {
    data: Vec<Pair<D, P>>,
    indices: HashMap<D, usize>,
}

impl<D: Hash + Eq + Clone + Send + Sync, P: PartialOrd + Clone + Send + Sync>
    IndexedBinaryHeap<D, P>
{
    pub fn new() -> Self {
        Self {
            data: Vec::new(),
            indices: HashMap::new(),
        }
    }

    fn sift_up(&mut self, i: usize) -> SiftResult {
        print!("{i}");
        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] {
                (self.data[i], self.data[parent_index]) =
                    (self.data[parent_index].clone(), self.data[i].clone());
                self.indices.insert(self.data[i].clone().data(), i).unwrap();
                self.indices
                    .insert(self.data[parent_index].clone().data(), parent_index)
                    .unwrap();
                self.sift_up(parent_index)
            } else {
                // Sift complete
                Ok(())
            }
        } else {
            // Tried to sift a non-existent index
            Err(SiftError::new(i, self.data.len()))
        }
    }

    fn sift_down(&mut self, i: usize) -> SiftResult {
        print!("{i}");
        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.indices.insert(
                        self.data[smaller_child_index].clone().data(),
                        smaller_child_index,
                    );
                    self.data[i] = smaller_child.clone();
                    self.indices.insert(smaller_child.data(), i);

                    // 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(())
            }
        }
    }

    fn delete_pair(&mut self, i: usize) -> Option<Pair<D, P>> {
        if i >= self.data.len() {
            return None;
        }
        let pair = self.data[i].clone();
        if i < self.data.len() - 1 {
            self.data[i] = self.data.pop().unwrap();
            if self.data[i] < pair {
                self.sift_up(i).unwrap();
            } else {
                self.sift_down(i).unwrap();
            }
        } else {
            self.data.pop().unwrap();
        }
        self.indices.remove(pair.get_data());
        Some(pair)
    }
}

impl<D: Hash + Eq + Clone + Send + Sync, P: PartialOrd + Clone + Send + Sync> FromIterator<(D, P)>
    for IndexedBinaryHeap<D, P>
{
    fn from_iter<T: IntoIterator<Item = (D, P)>>(iter: T) -> Self {
        let mut this = Self::new();
        for (i, (data, priority)) in iter.into_iter().enumerate() {
            if let Some(prev) = this.indices.insert(data.clone(), i) {
                this.indices.insert(data.clone(), prev).unwrap();
                this.data[prev] = Pair::new(data, priority);
            } else {
                this.data.push(Pair::new(data, priority));
            }
        }
        for i in (0..=(this.data.len() / 2)).rev() {
            this.sift_down(i).expect("Index error during heapify");
        }
        this
    }
}

impl<D: Hash + Eq + Clone + Send + Sync, P: PartialOrd + Clone + Send + Sync> IndexedBacking<D, P>
    for IndexedBinaryHeap<D, P>
{
    fn len(&self) -> usize {
        self.data.len()
    }

    fn contains(&self, data: &D) -> bool {
        self.indices.contains_key(data)
    }

    fn set(&mut self, data: D, priority: P) {
        if let Some(index) = self.indices.get(&data) {
            let pair = self.data.get_mut(*index).unwrap();
            let old_priority = pair.get_priority();
            if priority < *old_priority {
                pair.set_priority(priority);
                self.sift_up(*index).unwrap();
            } else {
                pair.set_priority(priority);
                self.sift_down(*index).unwrap();
            }
        } else {
            let final_index = self.data.len();
            self.indices.insert(data.clone(), final_index);
            self.data.push(Pair::new(data, priority));
            self.sift_up(final_index).unwrap();
        }
    }

    fn remove(&mut self, data: D) -> Option<P> {
        if let Some(index) = self.indices.get(&data) {
            if let Some(pair) = self.delete_pair(*index) {
                Some(pair.get_priority().clone())
            } else {
                None
            }
        } else {
            None
        }
    }

    fn pop(&mut self) -> Option<D> {
        if let Some(pair) = self.delete_pair(0) {
            Some(pair.data())
        } else {
            None
        }
    }
}