rust_adt/heap/
maximal_heap.rs

1pub struct MaximalHeap {
2    data: Vec<i32>,
3}
4
5impl MaximalHeap {
6    pub fn new() -> MaximalHeap {
7        MaximalHeap { data: Vec::new() }
8    }
9
10    pub fn insert(&mut self, d: i32) {
11        self.data.push(d);
12        self.bubble_up(self.data.len() - 1);
13    }
14
15    pub fn extact_max(&mut self) -> Option<i32> {
16        if self.data.len() == 0 {
17            return None;
18        };
19
20        let res = self.data[0];
21        if self.data.len() != 1 {
22            self.data[0] = self.data.pop().unwrap();
23            self.bubble_down(0);
24        } else {
25            self.data.pop();
26        }
27
28        Some(res)
29    }
30
31    fn bubble_down(&mut self, index: usize) {
32        let current = self.data[index];
33        let left = if index * 2 + 1 < self.data.len() {
34            Some(self.data[index * 2 + 1])
35        } else {
36            None
37        };
38        let right = if index * 2 + 2 < self.data.len() {
39            Some(self.data[index * 2 + 2])
40        } else {
41            None
42        };
43
44        if left == None {
45            return;
46        } else if right == None {
47            if left.unwrap() > current {
48                self.data[index * 2 + 1] = self.data[index];
49                self.data[index] = left.unwrap();
50            }
51        } else if left.unwrap() >= right.unwrap() {
52            if current < left.unwrap() {
53                self.data[index * 2 + 1] = current;
54                self.data[index] = left.unwrap();
55                self.bubble_down(index * 2 + 1); // bubble to child
56            }
57        } else {
58            if current < right.unwrap() {
59                self.data[index * 2 + 2] = current;
60                self.data[index] = right.unwrap();
61                self.bubble_down(index * 2 + 2); // bubble to child
62            }
63        }
64    }
65
66    fn bubble_up(&mut self, index: usize) {
67        let parent = self.data[index / 2];
68        let left = if (index / 2) * 2 + 1 < self.data.len() {
69            Some(self.data[(index / 2) * 2 + 1])
70        } else {
71            None
72        };
73        let right = if (index / 2) * 2 + 2 < self.data.len() {
74            Some(self.data[(index / 2) * 2 + 2])
75        } else {
76            None
77        };
78
79        if right != None {
80            let left = left.unwrap();
81            let right = right.unwrap();
82            if left >= parent && left > right {
83                self.data[index / 2] = left;
84                self.data[(index / 2) * 2 + 1] = parent;
85            } else if right > parent && left < right {
86                self.data[index / 2] = right;
87                self.data[(index / 2) * 2 + 2] = parent;
88            }
89        } else if left != None {
90            let left = left.unwrap();
91            if left >= parent {
92                self.data[index / 2] = left;
93                self.data[(index / 2) * 2 + 1] = parent;
94            }
95        }
96
97        if index != 0 {
98            self.bubble_up(index / 2);
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    #[test]
107    fn basic_inserts() {
108        let mut h = MaximalHeap::new();
109        assert_eq!(h.extact_max(), None);
110        h.insert(1);
111        h.insert(2);
112        h.insert(3);
113
114        assert_eq!(h.data, vec![3, 1, 2]);
115        assert_eq!(h.extact_max(), Some(3));
116        assert_eq!(h.data, vec![2, 1]);
117        assert_eq!(h.extact_max(), Some(2));
118        assert_eq!(h.data, vec![1]);
119        assert_eq!(h.extact_max(), Some(1));
120        assert_eq!(h.extact_max(), None);
121    }
122    #[test]
123    fn advanced_inserts() {
124        let mut h = MaximalHeap::new();
125        assert_eq!(h.extact_max(), None);
126        h.insert(6);
127        h.insert(1);
128        assert_eq!(h.data, vec![6, 1]);
129        h.insert(5);
130        assert_eq!(h.data, vec![6, 1, 5]);
131        h.insert(3);
132        assert_eq!(h.data, vec![6, 3, 5, 1]);
133        h.insert(10);
134        assert_eq!(h.data, vec![10, 6, 5, 1, 3]);
135        h.insert(4);
136        assert_eq!(h.data, vec![10, 6, 5, 1, 3, 4]);
137        h.insert(2);
138        assert_eq!(h.data, vec![10, 6, 5, 1, 3, 4, 2]);
139        h.insert(6);
140        assert_eq!(h.data, vec![10, 6, 5, 6, 3, 4, 2, 1]);
141
142        assert_eq!(h.extact_max(), Some(10));
143        assert_eq!(h.extact_max(), Some(6));
144        assert_eq!(h.extact_max(), Some(6));
145        assert_eq!(h.extact_max(), Some(5));
146        assert_eq!(h.extact_max(), Some(4));
147    }
148    #[test]
149    fn build_heap() {
150        assert_eq!(3 + 2, 5);
151    }
152}