告别重复造轮子:用Rust实现一个可大可小的通用“万能”二叉堆!在高性能编程领域,堆(Heap)是一个基石般的数据结构,广泛应用于优先队列、调度算法、以及各种高效的排序场景中。Rust语言标准库虽然提供了std::collections::BinaryHeap,但它是一个最大堆。本篇文
在高性能编程领域,堆(Heap) 是一个基石般的数据结构,广泛应用于优先队列、调度算法、以及各种高效的排序场景中。Rust 语言标准库虽然提供了 std::collections::BinaryHeap,但它是一个最大堆。
本篇文章将带你深入理解堆的原理,并利用 Rust 的泛型、函数指针和 Iterator Trait 的强大组合,亲手打造一个通用、灵活且高性能的二叉堆(Heap<T>)。它能够通过一个简单的比较函数切换,在最小堆(MinHeap) 和 最大堆(MaxHeap) 之间自由转换,让你彻底掌握这一核心算法结构。
实现一个 通用的二叉堆(Binary Heap)
/*
    heap
    implement a binary heap function
*/
use std::cmp::Ord;
use std::default::Default;
pub struct Heap<T>
where
    T: Default,
{
    count: usize,
    items: Vec<T>,
    comparator: fn(&T, &T) -> bool,
}
impl<T> Heap<T>
where
    T: Default,
{
    pub fn new(comparator: fn(&T, &T) -> bool) -> Self {
        Self {
            count: 0,
            items: vec![T::default()],
            comparator,
        }
    }
    pub fn len(&self) -> usize {
        self.count
    }
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    pub fn add(&mut self, value: T) {
        self.count += 1;
        // 1. 将新值添加到 items 的末尾(即 items[self.count])
        if self.items.len() > self.count {
            self.items[self.count] = value;
        } else {
            self.items.push(value);
        }
        // 2. 向上浮动 (Sift-Up)
        let mut current_idx = self.count;
        // 只要当前节点不是根节点 (idx > 1)
        while current_idx > 1 {
            let parent_idx = self.parent_idx(current_idx);
            // 检查当前节点和父节点是否违反堆属性
            // (即在 MinHeap 中,子节点比父节点小;在 MaxHeap 中,子节点比父节点大)
            if (self.comparator)(&self.items[current_idx], &self.items[parent_idx]) {
                // 违反属性,交换
                self.items.swap(current_idx, parent_idx);
                current_idx = parent_idx;
            } else {
                // 属性已满足,停止上浮
                break;
            }
        }
    }
    fn parent_idx(&self, idx: usize) -> usize {
        idx / 2
    }
    fn children_present(&self, idx: usize) -> bool {
        self.left_child_idx(idx) <= self.count
    }
    fn left_child_idx(&self, idx: usize) -> usize {
        idx * 2
    }
    fn right_child_idx(&self, idx: usize) -> usize {
        self.left_child_idx(idx) + 1
    }
    fn smallest_child_idx(&self, idx: usize) -> usize {
        let left_idx = self.left_child_idx(idx);
        let right_idx = self.right_child_idx(idx);
        // 1. 检查是否有右子节点
        if right_idx > self.count {
            // 只有左子节点(或没有子节点,但 children_present 已经保证了至少有左子节点)
            left_idx
        } else {
            // 2. 左右子节点都存在,使用比较器判断哪个更符合堆属性
            // (self.comparator)(a, b) 为 true,则 a 是我们想要的 (e.g., MinHeap 中 a 较小)
            if (self.comparator)(&self.items[left_idx], &self.items[right_idx]) {
                left_idx
            } else {
                right_idx
            }
        }
    }
}
impl<T> Heap<T>
where
    T: Default + Ord,
{
    /// Create a new MinHeap
    pub fn new_min() -> Self {
        Self::new(|a, b| a < b)
    }
    /// Create a new MaxHeap
    pub fn new_max() -> Self {
        Self::new(|a, b| a > b)
    }
}
impl<T> Iterator for Heap<T>
where
    T: Default,
{
    type Item = T;
    fn next(&mut self) -> Option<T> {
        if self.is_empty() {
            return None;
        }
        // 1. 交换根节点 (index 1) 和最后一个元素 (index self.count)
        self.items.swap(1, self.count);
        // 2. 弹出并返回旧的根节点(现在在末尾)
        self.count -= 1;
        // 因为 self.items[0] 是默认值,所以我们pop掉最后一个元素是安全的
        let extracted_value = self.items.pop().unwrap_or_default();
        // 3. 向下沉降 (Sift-Down)
        let mut current_idx = 1;
        while self.children_present(current_idx) {
            // 找到最符合堆属性的子节点索引
            let target_child_idx = self.smallest_child_idx(current_idx);
            // 检查当前节点是否违反堆属性与目标子节点进行比较
            // 如果子节点比当前节点更符合堆属性 (e.g., MinHeap 中子节点更小)
            if (self.comparator)(&self.items[target_child_idx], &self.items[current_idx]) {
                // 违反属性,交换
                self.items.swap(current_idx, target_child_idx);
                current_idx = target_child_idx;
            } else {
                // 属性已满足,停止下沉
                break;
            }
        }
        Some(extracted_value)
    }
}
pub struct MinHeap;
impl MinHeap {
    #[allow(clippy::new_ret_no_self)]
    pub fn new<T>() -> Heap<T>
    where
        T: Default + Ord,
    {
        Heap::new(|a, b| a < b)
    }
}
pub struct MaxHeap;
impl MaxHeap {
    #[allow(clippy::new_ret_no_self)]
    pub fn new<T>() -> Heap<T>
    where
        T: Default + Ord,
    {
        Heap::new(|a, b| a > b)
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_empty_heap() {
        let mut heap = MaxHeap::new::<i32>();
        assert_eq!(heap.next(), None);
    }
    #[test]
    fn test_min_heap() {
        let mut heap = MinHeap::new();
        heap.add(4);
        heap.add(2);
        heap.add(9);
        heap.add(11);
        assert_eq!(heap.len(), 4);
        assert_eq!(heap.next(), Some(2));
        assert_eq!(heap.next(), Some(4));
        assert_eq!(heap.next(), Some(9));
        heap.add(1);
        assert_eq!(heap.next(), Some(1));
    }
    #[test]
    fn test_max_heap() {
        let mut heap = MaxHeap::new();
        heap.add(4);
        heap.add(2);
        heap.add(9);
        heap.add(11);
        assert_eq!(heap.len(), 4);
        assert_eq!(heap.next(), Some(11));
        assert_eq!(heap.next(), Some(9));
        assert_eq!(heap.next(), Some(4));
        heap.add(1);
        assert_eq!(heap.next(), Some(2));
    }
}
这段 Rust 代码完整实现了一个 通用的二叉堆(Binary Heap),支持同时构建 最小堆 (MinHeap) 和 最大堆 (MaxHeap),并且通过泛型和函数指针实现了灵活的比较逻辑。
这段代码实现了一个通用的 堆 (Heap) 数据结构,它支持通过传入不同的比较函数构造 最小堆 或 最大堆。堆是一种完全二叉树,常用于 优先队列 (Priority Queue)、排序算法 (Heap Sort) 等场景。
pub struct Heap<T>
where
    T: Default,
{
    count: usize,                    // 当前堆中元素数量
    items: Vec<T>,                   // 存储堆元素的动态数组,下标从 1 开始
    comparator: fn(&T, &T) -> bool,  // 比较函数,用于控制堆的性质(大顶堆/小顶堆)
}items[0] 被保留为默认值(不参与运算),这使得计算父子节点索引更简单:
idx / 2idx * 2idx * 2 + 1comparator 是一个函数指针,它决定了堆的排序规则。
例如:
|a, b| a < b|a, b| a > bpub fn new(comparator: fn(&T, &T) -> bool) -> Self { ... }
pub fn len(&self) -> usize { self.count }
pub fn is_empty(&self) -> bool { self.count == 0 }这些函数提供了堆的基本管理功能。new 初始化堆并设置比较逻辑,len 和 is_empty 提供统计信息。
pub fn add(&mut self, value: T) { ... }插入新元素的逻辑是堆的核心:
将新元素插入数组末尾
self.items.push(value);
self.count += 1;向上比较并交换 从新插入的节点开始,不断和父节点比较。如果违反堆的性质(例如在最小堆中,子节点比父节点更小),就交换它们的位置。
while current_idx > 1 {
 let parent_idx = self.parent_idx(current_idx);
 if (self.comparator)(&self.items[current_idx], &self.items[parent_idx]) {
     self.items.swap(current_idx, parent_idx);
     current_idx = parent_idx;
 } else {
     break;
 }
}这一过程确保了堆在插入新元素后仍然保持有序。
实现了 Iterator trait 后,可以使用 .next() 来弹出堆顶元素。
fn next(&mut self) -> Option<T> {
    if self.is_empty() { return None; }
    // 1. 交换根节点和最后一个节点
    self.items.swap(1, self.count);
    // 2. 弹出堆顶
    let extracted_value = self.items.pop().unwrap_or_default();
    self.count -= 1;
    // 3. 从根开始下沉
    let mut current_idx = 1;
    while self.children_present(current_idx) {
        let target_child_idx = self.smallest_child_idx(current_idx);
        if (self.comparator)(&self.items[target_child_idx], &self.items[current_idx]) {
            self.items.swap(current_idx, target_child_idx);
            current_idx = target_child_idx;
        } else {
            break;
        }
    }
    Some(extracted_value)
}✅ next() 会不断返回当前堆顶元素,并在每次调用后重新维护堆的有序性。
这样我们可以像迭代器一样使用堆来进行排序(堆排序的核心机制)。
为方便用户使用,代码提供了两个结构体包装器:
pub struct MinHeap;
pub struct MaxHeap;它们的 new() 方法分别指定不同的比较逻辑:
pub fn new<T>() -> Heap<T>
where
    T: Default + Ord,
{
    Heap::new(|a, b| a < b)  // 最小堆
}
pub fn new<T>() -> Heap<T>
where
    T: Default + Ord,
{
    Heap::new(|a, b| a > b)  // 最大堆
}代码附带的测试验证了堆的基本功能:
#[test]
fn test_min_heap() {
    let mut heap = MinHeap::new();
    heap.add(4);
    heap.add(2);
    heap.add(9);
    assert_eq!(heap.next(), Some(2));
}这段代码展示了如何使用 Rust 泛型 + 函数指针 + Iterator trait 实现一个 灵活、安全、通用的二叉堆结构,支持最小堆与最大堆两种模式,并具备上浮、下沉、插入、弹出、迭代等核心功能。
这段 Rust 代码以工程实践的角度,完美地实现了通用二叉堆这一核心数据结构。其核心设计思想在于:
Heap<T> 泛型结构,并通过 T: Default 和 T: Ord 约束,确保代码的通用性和安全性。comparator: fn(&T, &T) -> bool,将堆的排序规则(大顶堆或小顶堆)与核心逻辑分离,实现了一套代码,两种功能。Iterator Trait 的 next() 方法,将堆顶元素的提取操作(pop)转化为标准的迭代行为,极大地增强了代码的 Rust 风格和可用性。add 方法中)和 下沉(Sift-Down)(next 方法中)两大核心操作,确保了堆的 $\mathcal{O}(\log n)$ 时间复杂度特性。总而言之,这个实现不仅是数据结构的学习范本,也是 Rust 语言高级特性(如 Trait、泛型、闭包与函数指针)的优秀实践案例。
 
                如果觉得我的文章对您有用,请随意打赏。你的支持将鼓励我继续创作!