fpq/
queue.rs

1use std::collections::BinaryHeap;
2
3use crate::scored_item::ScoredItem;
4
5/// A priority queue with a score function.
6pub struct PriorityQueue<T, S: Ord + Copy> {
7    heap: BinaryHeap<ScoredItem<T, S>>,
8    score_fn: Box<dyn Fn(&T) -> S>,
9}
10
11impl<T, S: Ord + Copy> PriorityQueue<T, S> {
12    /// Creates a new priority queue with the given score function.
13    pub fn new(score_fn: Box<dyn Fn(&T) -> S>) -> Self {
14        let heap = BinaryHeap::new();
15        Self { heap, score_fn }
16    }
17
18    /// Returns the reference to the first elements in the queue.
19    pub fn peek(&self) -> Option<(S, &T)> {
20        self.heap
21            .peek()
22            .map(|scored_item| (scored_item.score, &scored_item.item))
23    }
24
25    /// Returns the first elements in the queue.
26    pub fn pop(&mut self) -> Option<(S, T)> {
27        self.heap
28            .pop()
29            .map(|scored_item| (scored_item.score, scored_item.item))
30    }
31
32    /// Pushes an item into the queue.
33    /// The score of the item is calculated by the score function.
34    pub fn push(&mut self, item: T) {
35        let score = (self.score_fn)(&item);
36        self.push_with_score(item, score);
37    }
38
39    /// Pushes an item into the queue with the given score
40    /// without using the score function.
41    pub fn push_with_score(&mut self, item: T, score: S) {
42        self.heap.push(ScoredItem { item, score });
43    }
44
45    /// Converts the priority queue into a vector of (score, item) pairs.
46    pub fn into_vec(self) -> Vec<(S, T)> {
47        self.heap.into_iter().map(|x| (x.score, x.item)).collect()
48    }
49}
50
51#[cfg(test)]
52mod test {
53    use std::cmp::Reverse;
54
55    use super::PriorityQueue;
56    #[test]
57    fn test_priority_queue() {
58        let score_fn = Box::new(|s: &String| s.len());
59        let mut queue = PriorityQueue::new(score_fn);
60
61        assert!(queue.peek().is_none());
62
63        queue.push("a".to_string()); // score = 1
64        queue.push("ccc".to_string()); // score = 3
65        queue.push("bb".to_string()); // score = 2
66        queue.push_with_score("b".to_string(), 10); // score = 10
67
68        assert_eq!(queue.peek(), Some((10, &"b".to_string())));
69        assert_eq!(queue.pop(), Some((10, "b".to_string())));
70        assert_eq!(queue.pop(), Some((3, "ccc".to_string())));
71        assert_eq!(queue.pop(), Some((2, "bb".to_string())));
72        assert_eq!(queue.pop(), Some((1, "a".to_string())));
73        assert!(queue.peek().is_none());
74    }
75
76    #[test]
77    fn test_priority_queue_reverse_order() {
78        let score_fn = Box::new(|s: &String| Reverse(s.len()));
79        let mut queue = PriorityQueue::new(score_fn);
80
81        queue.push("a".to_string()); // score = -1
82        queue.push("ccc".to_string()); // score = -3
83        queue.push("bb".to_string()); // score = -2
84        queue.push_with_score("b".to_string(), Reverse(10)); // score = -10
85
86        assert_eq!(queue.peek(), Some((Reverse(1), &"a".to_string())));
87        assert_eq!(queue.pop(), Some((Reverse(1), "a".to_string())));
88        assert_eq!(queue.pop(), Some((Reverse(2), "bb".to_string())));
89        assert_eq!(queue.pop(), Some((Reverse(3), "ccc".to_string())));
90        assert_eq!(queue.pop(), Some((Reverse(10), "b".to_string())));
91        assert!(queue.peek().is_none());
92    }
93
94    #[test]
95    fn test_into_vec() {
96        let score_fn = Box::new(|s: &String| s.len());
97        let mut queue = PriorityQueue::new(score_fn);
98
99        queue.push("a".to_string()); // score = 1
100        queue.push("ccc".to_string()); // score = 3
101        queue.push("bb".to_string()); // score = 2
102        queue.push_with_score("b".to_string(), 10); // score = 10
103
104        let mut vec = queue.into_vec();
105        vec.sort_by_key(|x| x.0);
106        assert_eq!(
107            vec,
108            vec![
109                (1, "a".to_string()),
110                (2, "bb".to_string()),
111                (3, "ccc".to_string()),
112                (10, "b".to_string()),
113            ]
114        );
115    }
116}
fpq/queue.rs

fpq/
queue.rs
