rusty_priority_queue/
lib.rs

1use std::{fmt::Display, collections::HashMap};
2use std::hash::Hash;
3use anyhow::Result;
4use thiserror::Error;
5
6use pair::Pair;
7
8pub mod pair;
9
10#[derive(Error, Debug)]
11pub enum DHeapError {
12    #[error("Element already exists in the heap")]
13    ElementAlreadyExists,
14    #[error("usize max value is not available for priority")]
15    UnavailablePriority
16}
17
18
19#[derive(Debug)]
20pub struct DHeap<T: Eq + Hash + Clone + Display + PartialEq> {
21    data: Vec<Pair<T>>,
22    branching_factor: usize,
23    map: HashMap<T, bool>
24}
25
26impl<T: Eq + Hash + Clone + Display + PartialEq> DHeap<T> {
27    /// Creates a new heap
28    pub fn new(initial_capacity: Option<usize>, branching_factor: Option<usize>) -> Self {
29        match initial_capacity {
30            Some(v) => DHeap { data: Vec::with_capacity(v), 
31                branching_factor: branching_factor.unwrap_or(4),
32                map: HashMap::with_capacity(v)},
33            None => DHeap { data: Vec::new(),
34                branching_factor: branching_factor.unwrap_or(4),
35                map: HashMap::new()},
36        }
37    }
38
39    /// Accepts a slice of pairs and creates a heap
40    pub fn with_pairs(data: &[Pair<T>], initial_capacity: Option<usize>, branching_factor: Option<usize>) -> Result<Self> {
41        if data.iter().any(|x| x.priority == std::usize::MAX) { return Err(anyhow::Error::new(DHeapError::UnavailablePriority)); }
42
43        let capacity = if let Some(capacity) = initial_capacity {
44            if capacity > data.len() { capacity } else { data.len() * 2 }
45            } else { data.len() * 2 };
46        
47        let mut heap = DHeap { data: Vec::with_capacity(capacity), 
48                branching_factor: branching_factor.unwrap_or(4),
49                map: HashMap::with_capacity(capacity)};
50        heap.map = data.iter().map(|x| (x.get_cloned_element(), true)).collect::<HashMap<T, bool>>();
51        heap.data = Vec::from(data);
52        heap.heapify();
53            
54        Ok(heap)
55    }
56
57    /// Returns if the element exists in the heap
58    pub fn contains(&self, element: &T) -> bool {
59        if self.map.contains_key(element) { true }
60        else { false }
61    }
62
63    /// Removes the element from the heap
64    // Essentially we're just updating its priority to the max, then pop
65    // Due to that, we shouldn't allow max usize priority while inserting
66    pub fn remove(&mut self, element: T) -> Option<Pair<T>> {
67        if !self.map.contains_key(&element) { return None; }
68        self.map.remove(&element);
69
70        self.update_priority(element, std::usize::MAX);
71        self.top()
72    }
73
74    /// Inserts the value
75    pub fn insert_value(&mut self, element: T, priority: usize) -> Result<(), anyhow::Error> {
76        if self.map.contains_key(&element) { return Err(anyhow::Error::new(DHeapError::ElementAlreadyExists)); }
77        if priority == std::usize::MAX { return Err(anyhow::Error::new(DHeapError::UnavailablePriority)); }
78
79        self.map.insert(element.clone(), true);
80        
81        let pair = Pair::new(element, priority);
82        self.data.push(pair);
83        self.bubble_up(None);
84
85        Ok(())
86    }
87
88    /// Inserts a pair
89    pub fn insert_pair(&mut self, element: Pair<T>) -> Result<(), anyhow::Error> {
90        if self.map.contains_key(&element.get_element()) { return Err(anyhow::Error::new(DHeapError::ElementAlreadyExists)); }
91        if element.priority == std::usize::MAX { return Err(anyhow::Error::new(DHeapError::UnavailablePriority)); }
92
93        self.map.insert(element.get_cloned_element(), true);
94
95        self.data.push(element);
96        self.bubble_up(None);
97
98        Ok(())
99    }
100
101    /// Returns the highest priority value without taking it out of the queue
102    /// If empty, returns None
103    pub fn peek(&self) -> Option<&Pair<T>> {
104        if self.data.len() == 0 {
105            None
106        } else {
107            Some(&self.data[0])
108        }
109    }
110    
111    /// Returns the highest priority value. This operation take the value out of the queue
112    /// If empty, returns None
113    pub fn top(&mut self) -> Option<Pair<T>> {
114        let last_element = self.remove_last()?;
115        if self.data.is_empty() {
116            self.map.remove(last_element.get_element());
117            Some(last_element)
118        } else {
119            let root_element = self.data[0].clone();
120            self.data[0] = last_element;
121            self.push_down_optimized(None);
122            self.map.remove(root_element.get_element());
123            Some(root_element)
124        }
125    }
126
127    /// Finds and update priority of the value
128    pub fn update_priority(&mut self, old_value: T, new_priority: usize) {
129        if let Some(index) = self.find_index(old_value) {
130            let temp = self.data[index].clone();
131            self.data.remove(index);
132            let updated_pair = Pair::new(temp.get_cloned_element(), new_priority);
133            self.insert_pair_for_update(updated_pair);
134        }
135    }
136
137    fn insert_pair_for_update(&mut self, element: Pair<T>) {
138        self.data.push(element);
139        self.bubble_up(None);
140    }
141
142    fn heapify(&mut self)
143    {
144        let mut max_index = (self.data.len() - 1) / self.branching_factor;
145        while max_index != 0 {
146            self.push_down_optimized(Some(max_index));
147            max_index -= 1;
148        }
149        self.push_down_optimized(None);
150    }
151
152    fn find_index(&self, old_value: T) -> Option<usize> {
153        for (index, pair) in self.data.iter().enumerate() {
154            if *pair.get_element() == old_value {
155                return Some(index);
156            }
157        }
158        None
159    }
160
161    fn remove_last(&mut self) -> Option<Pair<T>> {
162        if self.data.is_empty() {
163            None
164        } else {
165            self.data.pop()
166        }
167    }
168
169    // bubbles up the selected element
170    fn bubble_up(&mut self, index: Option<usize>) {
171        // as default the last element is selected
172        let mut parent_index = index.unwrap_or(self.data.len() - 1);
173        while parent_index > 0 {
174            let current_index = parent_index;
175            parent_index = self.get_parent_index(parent_index);
176            if self.data[parent_index].priority < self.data[current_index].priority {
177                self.swap(current_index, parent_index)
178            } else {
179                break;
180            }
181        }
182    }
183
184    #[allow(dead_code)]
185    fn bubble_up_optimized(&mut self, initial_index: Option<usize>) {
186        let mut index = initial_index.unwrap_or(self.data.len() - 1);
187        let current = self.data[index].clone();
188        while index > 0 {
189            let parent_index = self.get_parent_index(index);
190            if self.data[parent_index].priority < self.data[index].priority {
191                self.data[index] = self.data[parent_index].clone();
192                index = parent_index;
193            } else {
194                break;
195            }
196        }
197        self.data[index] = current;
198    }
199
200    #[allow(dead_code)]
201    fn push_down(&mut self, initial_index: Option<usize>) {
202        let index = initial_index.unwrap_or(0);
203        let mut current_index = index;
204        while current_index < self.first_leaf_index() {
205            let highest_priority_child_index = self.highest_priority_child_index(index);
206            if self.data[current_index].priority < self.data[highest_priority_child_index].priority {
207                self.swap(current_index,highest_priority_child_index);
208                current_index = highest_priority_child_index;
209            } else {
210                break;
211            }
212        }    
213    }
214
215    fn push_down_optimized(&mut self, initial_index: Option<usize>) {
216        let mut index = initial_index.unwrap_or(0);
217        let current = self.data[index].clone();
218        while index < self.first_leaf_index() {
219            let highest_priority_child_index = self.highest_priority_child_index(index);
220            if self.data[index].priority < self.data[highest_priority_child_index].priority {
221                self.data[index] = self.data[highest_priority_child_index].clone();
222                index = highest_priority_child_index;
223            } else {
224                break;
225            }
226        } 
227        self.data[index] = current;
228    }
229
230    fn first_leaf_index(&self) -> usize {
231        (self.data.len() - 2) / self.branching_factor + 1
232    }
233
234    fn get_parent_index(&self, index: usize) -> usize {
235        (index - 1) / self.branching_factor
236    }
237
238    fn swap(&mut self, first_index: usize, second_index: usize) {
239        self.data.swap(first_index, second_index);
240    }
241
242    fn highest_priority_child_index(&self, index: usize) -> usize {
243        // if it has no child, returns itself
244        let first_child_index = (self.branching_factor * index) + 1;
245        if self.data.len() - 1 < first_child_index {
246            return index;
247        }
248
249        let mut highest_priority_index = index;
250        for i in 1..=self.branching_factor {
251            let child_index = (self.branching_factor * index) + i;
252            if self.data.len() - 1 < child_index {
253                continue;
254            }
255
256            if self.data[child_index].priority > self.data[highest_priority_index].priority {
257                highest_priority_index = child_index;
258            }
259        }
260        highest_priority_index
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    fn testing_dheap() -> DHeap<String> {
269        let mut heap = DHeap::new(None, None);
270        for i in 1..10 {
271            let example_pair = Pair::new(i.to_string(), i);
272            _ = heap.insert_pair(example_pair);
273        }
274        heap
275    }
276
277    #[test]
278    fn get_top_test() {
279        let mut heap = testing_dheap();
280        let received = heap.top().unwrap();
281        let expected = Pair::new("9", 9);
282
283        assert_eq!(expected.priority, received.priority);
284        assert_eq!(*expected.get_element(), *received.get_element());
285    }
286
287    #[test]
288    fn peek_test() {
289        let heap = testing_dheap();
290        let pair = heap.peek().unwrap();
291        assert_eq!(9, pair.priority);
292    }
293
294    #[test]
295    fn update_is_correct() {
296        let mut heap = testing_dheap();
297        _ = heap.update_priority("9".to_string(), 10);
298        let top_pair = heap.top().unwrap();
299        assert_eq!(10, top_pair.priority);
300    }
301
302    #[test] 
303    fn heapify_top_correct() {
304        let pairs = vec![Pair::new("9", 9), Pair::new("4", 4), Pair::new("11", 11),
305        Pair::new("10", 10), Pair::new("6", 6), Pair::new("20", 20)];
306        let mut heap = DHeap::with_pairs(&pairs, None, Some(4)).unwrap();
307        assert_eq!(20, heap.top().unwrap().priority);
308    }
309
310    #[test] 
311    fn remove_element() {
312        let pairs = vec![Pair::new("9", 9), Pair::new("4", 4), Pair::new("11", 11),
313        Pair::new("10", 10), Pair::new("6", 6), Pair::new("20", 20)];
314        let mut heap = DHeap::with_pairs(&pairs, None, Some(4)).unwrap();
315        assert_eq!("11", *heap.remove("11").unwrap().get_element());
316    }
317
318    #[test] 
319    fn contains_correct() {
320        let pairs = vec![Pair::new("9", 9), Pair::new("4", 4), Pair::new("11", 11),
321        Pair::new("10", 10), Pair::new("6", 6), Pair::new("20", 20)];
322        let heap = DHeap::with_pairs(&pairs, None, Some(4)).unwrap();
323        assert!(heap.contains(&"11"));
324    }
325}