gistools/data_structures/
priority_queue.rs

1use alloc::vec::Vec;
2
3/// Comparison function type for priority queue elements.
4pub type PriorityCompare<T> = fn(&T, &T) -> core::cmp::Ordering;
5
6/// # Priority Queue
7///
8/// ## Description
9/// A priority queue is a data structure that stores elements in a specific order.
10///
11/// A Rust port of the [tinyqueue](https://github.com/mourner/tinyqueue) code.
12///
13/// ## Usage
14/// ```rust
15/// use gistools::data_structures::PriorityQueue;
16///
17/// let mut queue = PriorityQueue::new(|a: &i32, b: &i32| a.cmp(b));
18///
19/// queue.push(3);
20/// queue.push(1);
21/// queue.push(2);
22///
23/// assert_eq!(queue.peek(), Some(&1));
24/// assert_eq!(queue.pop(), Some(1));
25/// assert_eq!(queue.len(), 2);
26/// ```
27///
28/// ## Links
29/// - <https://github.com/mourner/tinyqueue>
30#[derive(Debug)]
31pub struct PriorityQueue<T> {
32    data: Vec<T>,
33    compare: PriorityCompare<T>,
34}
35
36impl<T> PriorityQueue<T>
37where
38    T: Clone,
39{
40    /// Creates a new priority queue with an optional comparison function.
41    pub fn new(compare: PriorityCompare<T>) -> Self {
42        Self { data: Vec::new(), compare }
43    }
44
45    /// Returns the number of elements in the queue.
46    pub fn len(&self) -> usize {
47        self.data.len()
48    }
49
50    /// Returns true if the queue is empty.
51    pub fn is_empty(&self) -> bool {
52        self.data.is_empty()
53    }
54
55    /// Pushes an item into the queue.
56    pub fn push(&mut self, item: T) {
57        self.data.push(item);
58        self.up(self.data.len() - 1);
59    }
60
61    /// Removes and returns the highest-priority item.
62    pub fn pop(&mut self) -> Option<T> {
63        if self.data.is_empty() {
64            return None;
65        }
66        let top = self.data.swap_remove(0);
67        if !self.data.is_empty() {
68            self.down(0);
69        }
70        Some(top)
71    }
72
73    /// Peeks at the highest-priority item without removing it.
74    pub fn peek(&self) -> Option<&T> {
75        self.data.first()
76    }
77
78    fn up(&mut self, mut pos: usize) {
79        let item = self.data[pos].clone();
80        while pos > 0 {
81            let parent = (pos - 1) / 2;
82            if (self.compare)(&self.data[parent], &item) != core::cmp::Ordering::Greater {
83                break;
84            }
85            self.data[pos] = self.data[parent].clone();
86            pos = parent;
87        }
88        self.data[pos] = item;
89    }
90
91    fn down(&mut self, mut pos: usize) {
92        let len = self.data.len();
93        let item = self.data[pos].clone();
94        let half_len = len / 2;
95        while pos < half_len {
96            let mut child = 2 * pos + 1;
97            if child + 1 < len
98                && (self.compare)(&self.data[child + 1], &self.data[child])
99                    == core::cmp::Ordering::Less
100            {
101                child += 1;
102            }
103            if (self.compare)(&self.data[child], &item) != core::cmp::Ordering::Less {
104                break;
105            }
106            self.data[pos] = self.data[child].clone();
107            pos = child;
108        }
109        self.data[pos] = item;
110    }
111}