Rust——leetcode_剑指 Offer 40. 最小的k个数

使用最大堆的优先队列，解决最小的k个数问题
最大堆优先队列先存储前k个元素，遍历剩下的元素，如果元素比队首元素小，就入队，否则忽略
最终队列存储了最小的k个数

impl Solution {
    pub fn get_least_numbers(arr: Vec<i32>, k: i32) -> Vec<i32> {
        let mut queue = PriorityQueue::new_priority();
        for i in 0..k {
            queue.enqueue(arr[i as usize])
        }
        for i in k..arr.len() as i32 {
            if let Some(d) = queue.get_front() {
                if !queue.is_empty() && arr[i as usize] < d {
                    queue.dequeue();
                    queue.enqueue(arr[i as usize]);
                }
            }
        }
        return queue.max_heap.data.to_vec()
    }
}

trait Queue<E> {
    fn get_size(&self) -> usize;
    fn is_empty(&self) -> bool;
    fn enqueue(&mut self, e: E);
    fn dequeue(&mut self) -> Option<E>;
    fn get_front(&self) -> Option<E>;
}

struct PriorityQueue<E: PartialOrd + Copy> {
    max_heap: MaxHeap<E>
}

impl<E: PartialOrd + Copy> PriorityQueue<E> {
    fn new_priority() -> PriorityQueue<E> {
        Self {
            max_heap: MaxHeap::new_max_heap()
        }
    }
}

impl<E: PartialOrd + Copy> Queue<E> for PriorityQueue<E> {
    fn get_size(&self) -> usize {
        self.max_heap.get_size()
    }

    fn is_empty(&self) -> bool {
        self.max_heap.is_empty()
    }

    fn enqueue(&mut self, e: E) {
        self.max_heap.add(e)
    }

    fn dequeue(&mut self) -> Option<E> {
        self.max_heap.extract_max()
    }

    fn get_front(&self) -> Option<E> {
        self.max_heap.find_max()
    }
}

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;
        }
    }
}

执行结果