Rust——最大堆

实现

#![allow(unused)]

pub struct MaxHeap<E: PartialOrd + Copy> {
    data: Vec<E>
}

impl<E: PartialOrd + Copy> MaxHeap<E> {
    pub fn new_max_heap() -> MaxHeap<E> {
        Self {
            data: Vec::new()
        }
    }

    pub fn new_max_heap_with_capacity(capacity: usize) -> MaxHeap<E> {
        Self {
            data: Vec::with_capacity(capacity)
        }
    }

    pub fn get_size(&self) -> usize {
        self.data.len()
    }

    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }
    // index索引的父节点
    fn parent(&self, index: usize) -> Option<usize> {
        if index == 0 {
            return None;
        }
        Some((index - 1) / 2)
    }

    fn left_child(&self, index: usize) -> usize {
        index * 2 + 1
    }
    fn right_child(&self, index: usize) -> usize {
        index * 2 + 2
    }

    // 维持完全二叉树，将添加的元素放到树的末尾，上浮操作，以满足最大堆的性质
    pub fn add(&mut self, e: E) {
        self.data.push(e);
        self.sift_up(self.data.len() - 1)
    }
    // 上浮
    // 和父节点比较，大于父节点，就交换位置
    fn sift_up(&mut self, mut index: usize) {
        // 父节点小于字节点
        while index > 0 {
            let res = self.parent(index);
            if let Some(p_ind) = res {
                if self.data.get(p_ind) < self.data.get(index) {
                    // swap
                    self.data.swap(index, p_ind);
                    index = p_ind;
                } else {
                    break;
                }
            } else {
                return;
            }
        }
    }

    pub fn find_max(&self) -> Option<E> {
        if self.data.len() == 0 {
            return None;
        }
        Some(self.data[0])
    }
    // 取出最大值，就是取出第一个元素，为了维持完全二叉树，将最后一个叶子节点，替换到第一个元素位置
    // 然后下沉操作，以满足最大堆的性质
    pub fn extract_max(&mut self) -> Option<E> {
        let res = self.find_max();
        let end_ind = self.data.len() - 1;
        // 将末尾元素提到根
        self.data.swap(0, end_ind);
        // 删除最后一个元素
        self.data.remove(end_ind);
        // 下沉
        self.sift_down(0);
        res
    }
    // 下沉
    // 左右子树比较，和最大的交换位置
    fn sift_down(&mut self, mut index: usize) {
        let res = self.left_child(index);
        //
        while self.left_child(index) < self.data.len() {
            let mut j = self.left_child(index);
            // 有右子节点，右节点元素大
            if j + 1 < self.data.len() && self.data.get(j + 1) > self.data.get(j) {
                j += 1;
            }
            // 此时j是最大值索引
            if self.data.get(index) > self.data.get(j) {
                break;
            }
            self.data.swap(index, j);
            index = j;
        }
    }

    // replace
    pub fn replace(&mut self, e: E) -> Option<E> {
        let res = self.find_max();
        // 替换
        self.data.insert(0, e);
        // 下沉
        self.sift_down(0);
        res
    }

    // 把数组进行堆化，O(n)复杂度
    // 从最后一个非叶子节点进行下沉
    pub fn heapify(&mut self, data: Vec<E>) -> bool {
        self.data = data;
        if self.data.len() != 1 {
            // 最后一个非叶子节点，就是最后一个叶子节点的父节点
            let mut p_ind = self.parent(self.data.len() - 1);
            match p_ind {
                None => false,
                Some(ind) => {
                    let mut i = ind as i32;
                    while i >= 0 {
                        self.sift_down(i as usize);
                        i -= 1;
                    }
                    true
                }
            };
        }
        return true;
    }
}

排序

    // 根据传入的数组，转化最大堆，再排序
    pub fn sort(mut data: Vec<E>) -> MaxHeap<E> {
        if data.len() <= 1 {
            return Self { data };
        }
        // 最后一个叶子节点
        let index = data.len() - 1;
        // 父节点索引
        let p_ind = (index - 1) / 2;
        let mut i = p_ind as i32;
        while i >= 0 {
            let le = data.len();
            Self::sift_down_sort(&mut data, i as usize, le);
            i -= 1;
        }
        let mut j = (data.len() - 1) as i32;
        while j >= 0 {
            // 将最大值放到最后
            data.swap(0, j as usize);
            // 维持最大堆性质
            Self::sift_down_sort(&mut data, 0, j as usize);
            j -= 1;
        }
        Self {
            data
        }
    }
    // data[0,heap_ind)最大堆中，对索引index进行sift down
    fn sift_down_sort(data: &mut Vec<E>, mut index: usize, heap_ind: usize) {
        while 2 * index + 1 < heap_ind {
            let mut j = 2 * index + 1;
            // 有右子节点，右节点元素大
            if j + 1 < heap_ind && data.get(j + 1) > data.get(j) {
                j += 1;
            }
            // 此时j是最大值索引
            if data.get(index) > data.get(j) {
                break;
            }
            data.swap(index, j);
            index = j;
        }
    }

测试

mod max_heap;

#[cfg(test)]
mod tests {
    use super::max_heap::*;

    #[test]
    fn it_works() {
        let mut heap = MaxHeap::new_max_heap();
        assert_eq!(heap.is_empty(), true);
        heap.add(1);
        heap.add(3);
        heap.add(2);
        heap.add(8);
        assert_eq!(heap.extract_max(), Some(8));
        assert_eq!(heap.find_max(), Some(3));

        assert_eq!(heap.replace(10), Some(3));
        assert_eq!(heap.extract_max(), Some(10));
        let data = vec![1, 2, 3, 14, 5, 6];
        assert_eq!(heap.heapify(data), true);
        assert_eq!(heap.extract_max(), Some(14));

        let res = vec![1, 2, 3, 24, 5, 6];
        let mut data = MaxHeap::sort(res);
        assert_eq!(data.find_max(), Some(1));
    }
}