prio_queue/

lib.rs

////////////////////////////////////////////////////////////////////////////////
//! A simple priority queue implemented as a binary heap stored in a `Vec<T>`.
//!
//! This implementation aims to be simple and more flexible than the one
//! provided by the standard library. It relies on a given function to determine
//! priority between elements. Let you acces the underlying `Vec<T>` with
//! consistency check or not.

use std::fmt;

/// A priority queue implemented as a binary heap stored in a `Vec<T>`
///
/// This implementation aims to be simple and more flexible than the one
/// provided by the standard library. It rely on a given function to determine
/// priority between elements. Let you acces the underlying `Vec<T>` with
/// consistency check or not.
/// ```
/// # use prio_queue::PrioQueue;
/// let mut queue = PrioQueue::new(|l, r| l < r);
/// queue.push(42);
/// queue.push(32);
/// queue.push(64);
/// assert_eq!(queue.pop(), Some(32));
/// assert_eq!(queue.pop(), Some(42));
/// assert_eq!(queue.pop(), Some(64));
/// assert_eq!(queue.pop(), None);
/// ```
///
/// ## Complexity
///
/// Pushing an element has a maximum complexity of *O*(log(n)).
///
/// Poping en element has a maximum complexity of *O*(log(n)).
#[derive(Clone)]
pub struct PrioQueue<T> {
    prio: fn(&T, &T)->bool,
    tree: Vec<T>,
}
impl<T> PrioQueue<T> {

    /// Constructs a new, empty `PrioQueue<T>`.
    pub fn new(prio: fn(&T, &T)->bool) -> Self {
        Self {
            prio,
            tree: Vec::new(),
        }
    }

    pub fn with_capacity(capacity: usize, prio: fn(&T, &T)->bool) -> Self {
        Self {
            prio,
            tree: Vec::with_capacity(capacity),
        }
    }

    /// Convenient function that uses `|l, r| l < r` as the priority function.
    pub fn new_min() -> Self
    where T: PartialOrd
    {
        Self::new(|l, r| l < r)
    }

    /// Convenient function that uses `|l, r| l > r` as the priority function.
    pub fn new_max() -> Self
    where T: PartialOrd
    {
        Self::new(|l, r| l > r)
    }

    /// Inserts an element in the queue.
    pub fn push(&mut self, value: T) {
        let index = self.len();
        self.tree.push(value);
        self.bubble_up(index);
    }

    /// Removes the top element of the queue.
    pub fn pop(&mut self) -> Option<T> {
        if self.len() == 0 {
            None
        }
        else {
            let poped = self.tree.swap_remove(0);
            if self.len() > 0 {
                self.bubble_down(0);
            }
            Some(poped)
        }
    }

    /// Removes a given element.
    pub fn remove(&mut self, index: usize) -> T {
        assert!(index < self.len());
        let poped = self.tree.swap_remove(index);
        if index < self.len() {
            self.bubble_up(index);
            self.bubble_down(index);
        }
        poped
    }

    /// Returns a draining iterator yielding element in priority order.
    pub fn drain(&mut self) -> Drain<T> {
        Drain {
            inner: self
        }
    }

    /// Retains only the elements specified by the predicate.
    pub fn retain(&mut self, mut pred: impl FnMut(&T)->bool) {
        let mut index = 0;
        while index < self.len() {
            if pred(&self[index]) {
                self.bubble_up(index);
                index += 1;
            }
            else {
                self.tree.swap_remove(index);
            }
        }
    }

    fn bubble_up(&mut self, mut index: usize) {
        debug_assert!(index < self.len());
        while let Some(parent) = self.parent(index) {
            if self.is_prio(index, parent) {
                self.tree.swap(index, parent);
                index = parent;
            }
            else {
                break
            }
        }
    }

    fn bubble_down(&mut self, mut index: usize) {
        debug_assert!(index < self.len());
        while let Some((left, right)) = self.childs(index) {
            let result = if let Some(right) = right {
                if self.is_prio(left, right) {
                    if self.is_prio(left, index) {
                        Some(left)
                    }
                    else if self.is_prio(right, index) {
                        Some(right)
                    }
                    else {
                        None
                    }
                }
                else {
                    if self.is_prio(right, index) {
                        Some(right)
                    }
                    else if self.is_prio(left, index) {
                        Some(left)
                    }
                    else {
                        None
                    }
                }
            }
            else {
                if self.is_prio(left, index) {
                    Some(left)
                }
                else {
                    None
                }
            };
            if let Some(result) = result {
                self.tree.swap(result, index);
                index = result;
            }
            else {
                break;
            }
        }
    }

    /// Tests whether `lhs` has priority over `rhs`.
    pub fn is_prio(&self, lhs: usize, rhs: usize) -> bool {
        assert!(lhs < self.len());
        assert!(rhs < self.len());
        assert!(lhs != rhs);
        (self.prio)(&self[lhs], &self[rhs])
    }

    /// Returns the parent of the given element in the heap.
    pub fn parent(&self, index: usize) -> Option<usize> {
        assert!(index < self.len());
        match index {
            0 => None,
            n => Some((n + 1) / 2 - 1),
        }
    }

    /// Returns the childs of the given element in the heap.
    pub fn childs(&self, index: usize) -> Option<(usize, Option<usize>)> {
        let right = (index + 1) * 2;
        let left = right - 1;
        match (left < self.len(), right < self.len()) {
            ( true,  true) => Some((left, Some(right))),
            ( true, false) => Some((left, None)),
            (false,     _) => None,
        }
    }

    /// Let you access modify the underlying `Vec<T>` without checking consistency.
    pub fn unchecked_mut(&mut self) -> &mut Vec<T> {
        &mut self.tree
    }

    /// Let you access modify the underlying `Vec<T>` and will check consistency.
    pub fn checked_mut(&mut self) -> CheckedMut<T> {
        CheckedMut{ inner: self }
    }

    /// Tests whether the heap is in a consistent state.
    pub fn check_consistency(&self) -> bool {
        (1..self.len()).all(|i| self.is_prio(self.parent(i).unwrap(), i))
    }

    /// Repair heap consistency if broken.
    pub fn restore_consistency(&mut self) {
        for i in 1..self.len() {
            self.bubble_up(i);
        }
    }

    ////////////////////////////////////////////////////////////////////////////////
    /// Builds a heap using the given `vec`.
    pub fn from_vec(vec: Vec<T>, prio: fn(&T, &T)->bool) -> Self {
        let mut queue = Self {
            prio,
            tree: vec,
        };
        queue.restore_consistency();
        queue
    }

    /// Use the given `vec` as the heap and checking its consistency.
    pub fn from_vec_checked(vec: Vec<T>, prio: fn(&T, &T)->bool) -> Self {
        let queue = Self {
            prio,
            tree: vec,
        };
        assert!(queue.check_consistency());
        queue
    }

    /// Use the given `vec` as the heap without checking its consistency.
    pub fn from_vec_unchecked(vec: Vec<T>, prio: fn(&T, &T)->bool) -> Self {
        Self {
            prio,
            tree: vec,
        }
    }

    /// Buils the queue from an iterator.
    pub fn from_iter<I: IntoIterator<Item=T>>(iter: I, prio: fn(&T, &T)->bool) -> Self {
        Self::from_vec(iter.into_iter().collect(), prio)
    }
}

pub struct CheckedMut<'a, T> {
    inner: &'a mut PrioQueue<T>,
}
impl<'a, T> std::ops::Deref for CheckedMut<'a, T> {
    type Target = Vec<T>;
    fn deref(&self) -> &Self::Target {
        &self.inner.tree
    }
}
impl<'a, T> std::ops::DerefMut for CheckedMut<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner.tree
    }
}
impl<'a, T> Drop for CheckedMut<'a, T> {
    fn drop(&mut self) {
        assert!(self.inner.check_consistency());
    }
}

impl<T: fmt::Debug> fmt::Debug for PrioQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        self.tree.fmt(f)
    }
}
impl<T: fmt::Display> fmt::Display for PrioQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        let width = f.width().unwrap_or(3);
        let len = (self.len() + 1).next_power_of_two();
        let pow = len.trailing_zeros() as usize;
        let wl = (width - 1) / 2;
        let wr = (width - 1) - wl;
        for i in 0..pow {
            let span = (1 << (pow - 1 - i)) - 1;
            for j in 0..(1 << i) {
                let elem = (1 << i) - 1 + j;
                let childs = self.childs(elem);
                match childs {
                    None => {
                        write!(f, "{: >1$}", "", width * span)?;
                        if elem < self.len() {
                            write!(f, "{: ^1$}", self[elem], width)?;
                        }
                        else {
                            write!(f, "{: >1$}", "", width)?;
                        }
                        write!(f, "{: >1$}", "", width * span)?;
                    }
                    Some((_, right)) => {
                        write!(f, "{: >1$}", "", width * (span / 2) + wl)?;
                        write!(f, "╭")?;
                        write!(f, "{:─>1$}", "", width * (span / 2) + wr)?;
                        write!(f, "{:─^1$}", self[elem], width)?;
                        if let Some(_) = right {
                            write!(f, "{:─>1$}", "", width * (span / 2) + wl)?;
                            write!(f, "╮")?;
                            write!(f, "{: >1$}", "", width * (span / 2) + wr)?;
                        }
                        else {
                            write!(f, "{: >1$}", "", width * span)?;
                        }
                    }
                }
                write!(f, "{: >1$}", "", width)?;
            }
            writeln!(f, )?;
        }
        Ok(())
    }
}

impl<T> std::ops::Deref for PrioQueue<T> {
    type Target = [T];
    fn deref(&self) -> &Self::Target {
        &self.tree
    }
}

pub struct Drain<'a, T> {
    inner: &'a mut PrioQueue<T>,
}
impl<'a, T> Iterator for Drain<'a, T> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        self.inner.pop()
    }
}

pub struct IntoIter<T> {
    inner: PrioQueue<T>,
}
impl<T> Iterator for IntoIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        self.inner.pop()
    }
}

impl<T> IntoIterator for PrioQueue<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;
    fn into_iter(self) -> Self::IntoIter {
        IntoIter{
            inner: self
        }
    }
}

impl<T> Extend<T> for PrioQueue<T> {
    fn extend<I>(&mut self, iter: I)
    where I: IntoIterator<Item = T>
    {
        for elem in iter {
            self.push(elem);
        }
    }
}

impl<T: PartialOrd> std::iter::FromIterator<T> for PrioQueue<T> {
    fn from_iter<I>(iter: I) -> Self
    where I: IntoIterator<Item = T>
    {
        let mut queue = Self::new(|l, r| l < r);
        queue.extend(iter);
        queue
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn sorts() {
        use rand::{Rng, SeedableRng, rngs::SmallRng};
        let mut rng = SmallRng::seed_from_u64(9320154);
        for _i in 0..100 {
            let len = rng.gen_range(0..128);
            // println!("round {}/10: len = {}", i + 1, len);
            let mut values = Vec::with_capacity(len);
            for _ in 0..len {
                values.push(rng.gen_range(0..100));
            }
            let queue: PrioQueue<_> = values.iter().copied().collect();
            // print!("{:2}", queue);
            let result: Vec<_> = queue.into_iter().collect();
            values.sort();
            assert_eq!(values, result);
        }
    }
    #[test]
    fn remove() {
        use rand::{Rng, SeedableRng, rngs::SmallRng};
        let mut rng = SmallRng::seed_from_u64(885301);
        for _i in 0..100 {
            let len = rng.gen_range(0..128);
            let remove = rng.gen_range(0..=len);
            // println!("round {}/10: len = {}, remove = {}", i + 1, len, remove);
            let mut values = Vec::with_capacity(len);
            for _ in 0..len {
                values.push(rng.gen_range(0..100));
            }
            let mut queue: PrioQueue<_> = values.iter().copied().collect();
            // print!("{:2}", queue);
            for _ in 0..remove {
                let index = rng.gen_range(0..values.len());
                let value = values.swap_remove(index);
                let index = queue.iter().position(|&e| e == value).unwrap();
                let removed = queue.remove(index);
                assert_eq!(value, removed);
            }
            // print!("{:2}", queue);
            let result: Vec<_> = queue.into_iter().collect();
            values.sort();
            assert_eq!(values, result);
        }
    }

    #[test]
    fn retain() {
        use rand::{Rng, SeedableRng, rngs::SmallRng};
        let mut rng = SmallRng::seed_from_u64(340712856);
        for _i in 0..100 {
            let len = rng.gen_range(0..128);
            let factor = rng.gen_range(1..6);
            // println!("round {}/10: len = {}, factor = {}", i + 1, len, factor);
            let mut values = Vec::with_capacity(len);
            for _ in 0..len {
                values.push(rng.gen_range(0..100));
            }
            let mut queue: PrioQueue<_> = values.iter().copied().collect();
            // print!("{:2}", queue);
            queue.retain(|&e| e % factor != 0);
            values.retain(|&e| e % factor != 0);
            // print!("{:2}", queue);
            let result: Vec<_> = queue.into_iter().collect();
            values.sort();
            assert_eq!(values, result);
        }
    }

    #[test]
    #[should_panic]
    fn checked_bad_mut() {
        let mut queue = PrioQueue::new(|l, r| l < r);
        queue.push(1);
        queue.checked_mut().push(0);
    }

    #[test]
    fn unchecked_bad_mut() {
        use rand::{Rng, SeedableRng, rngs::SmallRng};
        use std::ops::RangeInclusive;

        fn test(len: usize, values: RangeInclusive<i32>, remove: usize, mut rng: impl Rng) {
            assert!(remove <= len);
            let mut vec = Vec::with_capacity(len);
            for _ in 0..len {
                vec.push(rng.gen_range(values.clone()));
            }
            let mut queue: PrioQueue<_> = vec.iter().copied().collect();
            let mut removed = {
                let tmp = queue.unchecked_mut();
                let mut removed = Vec::with_capacity(remove);
                for _ in 0..remove {
                    removed.push(tmp.remove(rng.gen_range(0..tmp.len())));
                }
                removed
            };
            assert_eq!(removed.len(), remove);
            assert_eq!(queue.len(), len - remove);
            removed.extend(queue);
            removed.sort();
            vec.sort();
            assert_eq!(removed, vec);
        }
        let mut rng = SmallRng::seed_from_u64(150738);
        for _i in 0..100 {
            let len = rng.gen_range(0..100);
            test(len, 0..=rng.gen_range(0..100), rng.gen_range(0..=len), &mut rng);
        }
    }
}