1use std::cmp::Ordering;
2
3use collection::{Collection, Disposable};
4
5use crate::traits::QueueLike;
6
7pub use crate::traits::PriorityQueueLike;
8
9pub struct PriorityQueue<T, C>
10where
11 C: Fn(&T, &T) -> Ordering,
12{
13 disposed: bool,
14 compare: C,
15 elements: Vec<T>,
16}
17
18impl<T, C> PriorityQueue<T, C>
19where
20 C: Fn(&T, &T) -> Ordering,
21{
22 pub fn new(compare: C) -> Self {
23 Self {
24 disposed: false,
25 compare,
26 elements: Vec::new(),
27 }
28 }
29
30 fn up(&mut self, index: usize) {
31 let mut q = index;
32 while q > 0 {
33 let p = (q - 1) >> 1;
34 if (self.compare)(&self.elements[p], &self.elements[q]) != Ordering::Greater {
35 break;
36 }
37 self.elements.swap(p, q);
38 q = p;
39 }
40 }
41
42 fn down(&mut self, index: usize) {
43 let mut p = index;
44 let n = self.elements.len();
45
46 while p < n {
47 let left = (p << 1) + 1;
48 if left >= n {
49 break;
50 }
51
52 let right = left + 1;
53 let mut q = left;
54 if right < n
55 && (self.compare)(&self.elements[right], &self.elements[left]) == Ordering::Less
56 {
57 q = right;
58 }
59
60 if (self.compare)(&self.elements[p], &self.elements[q]) != Ordering::Greater {
61 break;
62 }
63
64 self.elements.swap(p, q);
65 p = q;
66 }
67 }
68
69 fn fast_build(&mut self) {
70 let n = self.elements.len();
71 if n <= 1 {
72 return;
73 }
74
75 let last_parent = (n >> 1) - 1;
76 for p in (0..=last_parent).rev() {
77 self.down(p);
78 }
79 }
80}
81
82impl<T, C> QueueLike<T> for PriorityQueue<T, C>
83where
84 C: Fn(&T, &T) -> Ordering,
85{
86 fn front(&self) -> Option<&T> {
87 self.elements.first()
88 }
89
90 fn enqueue(&mut self, element: T) {
91 self.elements.push(element);
92 let index = self.elements.len() - 1;
93 self.up(index);
94 }
95
96 fn dequeue(&mut self) -> Option<T> {
97 if self.elements.is_empty() {
98 return None;
99 }
100 if self.elements.len() == 1 {
101 return self.elements.pop();
102 }
103
104 let removed = self.elements.swap_remove(0);
105 self.down(0);
106 Some(removed)
107 }
108
109 fn enqueues<I>(&mut self, elements: I)
110 where
111 I: IntoIterator<Item = T>,
112 {
113 let size = self.elements.len();
114 self.elements.extend(elements);
115
116 let next_size = self.elements.len();
117 if next_size == size {
118 return;
119 }
120
121 let new_added = next_size - size;
122 let next_size_f64 = next_size as f64;
123 if (new_added as f64) * next_size_f64.log2() > next_size_f64 {
124 self.fast_build();
125 } else {
126 for i in size..next_size {
127 self.up(i);
128 }
129 }
130 }
131
132 fn replace_front(&mut self, new_back: T) -> Option<T> {
133 if self.elements.is_empty() {
134 self.elements.push(new_back);
135 return None;
136 }
137
138 let removed = std::mem::replace(&mut self.elements[0], new_back);
139 self.down(0);
140 Some(removed)
141 }
142}
143
144impl<T, C> PriorityQueueLike<T> for PriorityQueue<T, C>
145where
146 C: Fn(&T, &T) -> Ordering,
147{
148 fn compare(&self, a: &T, b: &T) -> Ordering {
149 (self.compare)(a, b)
150 }
151}
152
153impl<T, C> Collection for PriorityQueue<T, C>
154where
155 C: Fn(&T, &T) -> Ordering,
156{
157 type Item = T;
158 type Iter<'a>
159 = std::slice::Iter<'a, T>
160 where
161 Self: 'a;
162
163 fn iter(&self) -> Self::Iter<'_> {
164 self.elements.iter()
165 }
166
167 fn size(&self) -> usize {
168 self.elements.len()
169 }
170
171 fn clear(&mut self) {
172 self.elements.clear();
173 }
174
175 fn retain<F>(&mut self, mut f: F) -> usize
176 where
177 F: FnMut(&Self::Item) -> bool,
178 {
179 let before = self.elements.len();
180 if before == 0 {
181 return 0;
182 }
183
184 self.elements.retain(|item| f(item));
185 let removed = before - self.elements.len();
186 if removed > 0 {
187 self.fast_build();
188 }
189 removed
190 }
191}
192
193impl<T, C> Disposable for PriorityQueue<T, C>
194where
195 C: Fn(&T, &T) -> Ordering,
196{
197 fn dispose(&mut self) {
198 self.disposed = true;
199 self.elements.clear();
200 }
201
202 fn is_disposed(&self) -> bool {
203 self.disposed
204 }
205}
206
207impl<'a, T, C> IntoIterator for &'a PriorityQueue<T, C>
208where
209 C: Fn(&T, &T) -> Ordering,
210{
211 type Item = &'a T;
212 type IntoIter = std::slice::Iter<'a, T>;
213
214 fn into_iter(self) -> Self::IntoIter {
215 self.elements.iter()
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use std::cmp::Ordering;
222
223 use collection::{Collection, Disposable};
224
225 use crate::traits::{PriorityQueueLike, QueueLike};
226
227 use super::PriorityQueue;
228
229 fn min_compare(a: &i32, b: &i32) -> Ordering {
230 a.cmp(b)
231 }
232
233 fn max_compare(a: &i32, b: &i32) -> Ordering {
234 b.cmp(a)
235 }
236
237 fn drain_all<C>(q: &mut PriorityQueue<i32, C>) -> Vec<i32>
238 where
239 C: Fn(&i32, &i32) -> Ordering,
240 {
241 let mut out = Vec::new();
242 while let Some(x) = q.dequeue() {
243 out.push(x);
244 }
245 out
246 }
247
248 #[test]
249 fn queue_like_min_heap_ops_should_work() {
250 let mut q = PriorityQueue::new(min_compare);
251
252 assert_eq!(q.front(), None);
253 assert_eq!(q.dequeue(), None);
254
255 q.enqueues([4, 2, 5, 1, 3]);
256 assert_eq!(q.front(), Some(&1));
257 assert_eq!(q.replace_front(6), Some(1));
258 assert_eq!(drain_all(&mut q), vec![2, 3, 4, 5, 6]);
259 }
260
261 #[test]
262 fn enqueue_should_cover_up_break_path() {
263 let mut q = PriorityQueue::new(min_compare);
264
265 q.enqueue(1);
266 q.enqueue(2);
267 q.enqueue(0);
268
269 assert_eq!(q.front(), Some(&0));
270 assert_eq!(drain_all(&mut q), vec![0, 1, 2]);
271 }
272
273 #[test]
274 fn replace_front_should_handle_empty_and_single_item() {
275 let mut q = PriorityQueue::new(min_compare);
276
277 assert_eq!(q.replace_front(10), None);
278 assert_eq!(q.front(), Some(&10));
279
280 assert_eq!(q.replace_front(5), Some(10));
281 assert_eq!(q.front(), Some(&5));
282 assert_eq!(q.dequeue(), Some(5));
283 }
284
285 #[test]
286 fn enqueues_should_work_for_small_and_large_batch() {
287 let mut q = PriorityQueue::new(min_compare);
288
289 q.enqueues([5, 4]);
290 q.enqueues(0..100);
291
292 assert_eq!(q.size(), 102);
293 assert_eq!(q.front(), Some(&0));
294
295 let drained = drain_all(&mut q);
296 assert_eq!(drained.len(), 102);
297 assert!(drained.windows(2).all(|w| w[0] <= w[1]));
298 }
299
300 #[test]
301 fn enqueues_empty_should_be_noop() {
302 let mut q = PriorityQueue::new(min_compare);
303
304 q.enqueues(std::iter::empty());
305 assert!(q.is_empty());
306
307 q.enqueue(3);
308 q.enqueues(std::iter::empty());
309 assert_eq!(q.front(), Some(&3));
310 assert_eq!(q.size(), 1);
311 }
312
313 #[test]
314 fn custom_comparator_should_support_max_heap() {
315 let mut q = PriorityQueue::new(max_compare);
316
317 q.enqueues([1, 5, 2, 4, 3]);
318
319 assert_eq!(q.front(), Some(&5));
320 assert_eq!(drain_all(&mut q), vec![5, 4, 3, 2, 1]);
321 }
322
323 #[test]
324 fn retain_should_rebuild_heap() {
325 let mut q = PriorityQueue::new(min_compare);
326 q.enqueues(1..=8);
327
328 let removed = q.retain(|x| *x % 2 == 0);
329 assert_eq!(removed, 4);
330 assert_eq!(drain_all(&mut q), vec![2, 4, 6, 8]);
331
332 let mut single = PriorityQueue::new(min_compare);
333 single.enqueues([1, 2]);
334 let removed_single = single.retain(|x| *x == 2);
335 assert_eq!(removed_single, 1);
336 assert_eq!(single.front(), Some(&2));
337 }
338
339 #[test]
340 fn retain_on_empty_should_return_zero() {
341 let mut q = PriorityQueue::new(min_compare);
342
343 let removed = q.retain(|_| true);
344 assert_eq!(removed, 0);
345 assert!(q.is_empty());
346 }
347
348 #[test]
349 fn iter_should_be_unsorted_but_complete() {
350 let mut q = PriorityQueue::new(min_compare);
351 q.enqueues([7, 1, 9, 3, 5]);
352
353 let mut from_iter: Vec<i32> = q.iter().copied().collect();
354 from_iter.sort();
355 assert_eq!(from_iter, vec![1, 3, 5, 7, 9]);
356
357 let mut from_into_iter: Vec<i32> = (&q).into_iter().copied().collect();
358 from_into_iter.sort();
359 assert_eq!(from_into_iter, vec![1, 3, 5, 7, 9]);
360 }
361
362 #[test]
363 fn collection_and_dispose_contract_should_work() {
364 let mut q = PriorityQueue::new(min_compare);
365 q.enqueues([3, 1, 2]);
366
367 assert_eq!(Collection::size(&q), 3);
368 assert_eq!(Collection::count(&q, |x| *x % 2 == 1), 2);
369
370 let mut all = Collection::collect(&q);
371 all.sort();
372 assert_eq!(all, vec![1, 2, 3]);
373
374 Collection::clear(&mut q);
375 assert!(Collection::is_empty(&q));
376
377 assert!(!Disposable::is_disposed(&q));
378 Disposable::dispose(&mut q);
379 assert!(Disposable::is_disposed(&q));
380 assert!(Collection::is_empty(&q));
381 }
382
383 #[test]
384 fn priority_queue_like_should_be_implemented() {
385 fn assert_priority_queue_like<Q: PriorityQueueLike<i32>>(_q: &Q) {}
386
387 let q = PriorityQueue::new(min_compare);
388 assert_priority_queue_like(&q);
389 assert_eq!(q.compare(&1, &2), Ordering::Less);
390 assert_eq!(q.compare(&2, &1), Ordering::Greater);
391 assert_eq!(q.compare(&3, &3), Ordering::Equal);
392 }
393}