使用 Rust 实现 Heap

1// MinHeap struct，它有一个 Vec<i32> 类型的字段，它是堆的内部存储
2pub struct MinHeap {
3    data: Vec<i32>,
4}
5
6#[allow(dead_code)]
7impl MinHeap {
8    // 新建一个空的 MinHeap
9    pub fn new() -> Self {
10        MinHeap { data: vec![] }
11    }
12
13    // 在堆中插入一个元素，然后进行上浮操作保持堆的性质。
14    // 我们将新元素添加到数组的末尾以维持完全二叉树的特性，然后进行上浮操作以维持堆的特性。
15    pub fn push(&mut self, item: i32) {
16        self.data.push(item);
17        self.bubble_up(self.data.len() - 1);
18    }
19
20    // 从堆中弹出最小元素。如果堆为空，返回 None。
21    // 否则，我们将堆的最后一个元素交换到堆顶（这保持了完全二叉树的特性），然后进行下沉操作以维持堆的特性。
22    pub fn pop(&mut self) -> Option<i32> {
23        if self.data.is_empty() {
24            return None;
25        }
26
27        let min = self.data[0];
28        let last = self.data.pop().unwrap();
29        if !self.data.is_empty() {
30            self.data[0] = last;
31            self.bubble_down(0);
32        }
33        Some(min)
34    }
35
36    // 上浮操作。从当前节点开始，如果当前节点的值小于其父节点的值，则交换这两个节点。
37    // 然后递归地对父节点进行上浮操作，直到当前节点的值不再小于其父节点，或者已经到达根节点。
38    // 这个操作用于维护插入新元素后的堆的性质。
39    fn bubble_up(&mut self, idx: usize) {
40        if idx > 0 {
41            let parent = (idx - 1) / 2;
42            if self.data[parent] > self.data[idx] {
43                self.data.swap(parent, idx);
44                self.bubble_up(parent);
45            }
46        }
47    }
48
49    // 下沉操作。从当前节点开始，如果当前节点的值大于其任一子节点的值，则交换这两个节点。
50    // 然后递归地对新的子节点进行下沉操作，直到当前节点的值不再大于其子节点，或者已经到达叶子节点。
51    // 这个操作用于维护移除最小元素后的堆的性质。
52    fn bubble_down(&mut self, idx: usize) {
53        let left = 2 * idx + 1;
54        let right = 2 * idx + 2;
55        let mut min_idx = idx;
56
57        // 找到当前节点，左子节点和右子节点中最小的那个
58        if left < self.data.len() && self.data[left] < self.data[min_idx] {
59            min_idx = left;
60        }
61        if right < self.data.len() && self.data[right] < self.data[min_idx] {
62            min_idx = right;
63        }
64
65        // 如果当前节点不是最小的，那么它就需要下沉
66        if min_idx != idx {
67            self.data.swap(min_idx, idx);
68            self.bubble_down(min_idx);
69        }
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn test_heap(){
79        let mut heap = MinHeap::new();
80        heap.push(1);
81        heap.push(4);
82        heap.push(5);
83        heap.push(2);
84        heap.push(3);
85        assert_eq!(heap.pop(), Some(1));
86        assert_eq!(heap.pop(), Some(2));
87        assert_eq!(heap.pop(), Some(3));
88        assert_eq!(heap.pop(), Some(4));
89        assert_eq!(heap.pop(), Some(5));
90        assert_eq!(heap.pop(), None);
91    }
92}