算法解惑：堆优先队列实战之TopK和3D接雨水

TopK-3D接雨水：堆优先队列进阶

#

#

#

引言

在数据结构与算法的世界里，堆和优先队列可谓是两大法宝，它们凭借其高效的组织方式和排序特性，在解决一系列实际问题时发挥着不可或缺的作用。本文将通过两个经典算法案例（TopK和3D接雨水）来深入浅出地解析堆和优先队列的原理和应用，并用C语言、Js和Rust三种语言进行具体实现，为读者提供全方位的理解和实践经验。

堆与优先队列

堆是一种完全二叉树，其每个节点都满足堆的性质：根节点的值总是比它的孩子节点的值大（小）。堆可以通过多种方式实现，其中最常见的是二叉堆和斐波那契堆。

优先队列是一种抽象数据类型，它支持以下操作：

• 入队：将一个元素插入优先队列中。
• 出队：从优先队列中删除并返回优先级最高的元素。
• 查看队首元素：返回优先队列中优先级最高的元素，但不删除它。

堆可以通过实现优先队列的操作来实现优先队列。

TopK问题

TopK问题是指从一个无序数组中找到最大的前K个元素。堆提供了一种有效率的解决方案。我们可以使用一个大小为K的最小堆，将数组中的元素逐个插入堆中。当堆中元素的个数超过K时，我们将堆顶元素（最小元素）弹出，并继续插入数组中的下一个元素。当数组中的所有元素都处理完毕后，堆中剩下的就是最大的前K个元素。

C语言实现：

#include <stdio.h>
#include <stdlib.h>

typedef struct heap {
    int *arr;
    int size;
    int capacity;
} heap_t;

heap_t *heap_create(int capacity) {
    heap_t *heap = malloc(sizeof(heap_t));
    heap->arr = malloc(sizeof(int) * capacity);
    heap->size = 0;
    heap->capacity = capacity;
    return heap;
}

void heap_insert(heap_t *heap, int val) {
    int i;
    if (heap->size == heap->capacity) {
        heap->capacity *= 2;
        heap->arr = realloc(heap->arr, sizeof(int) * heap->capacity);
    }
    i = heap->size++;
    heap->arr[i] = val;
    while (i > 0 && heap->arr[i] > heap->arr[(i - 1) / 2]) {
        int tmp = heap->arr[i];
        heap->arr[i] = heap->arr[(i - 1) / 2];
        heap->arr[(i - 1) / 2] = tmp;
        i = (i - 1) / 2;
    }
}

int heap_pop(heap_t *heap) {
    int val;
    if (heap->size == 0) {
        return -1;
    }
    val = heap->arr[0];
    heap->arr[0] = heap->arr[--heap->size];
    heapify(heap, 0);
    return val;
}

void heapify(heap_t *heap, int i) {
    int largest = i;
    int left = 2 * i + 1;
    int right = 2 * i + 2;
    if (left < heap->size && heap->arr[left] > heap->arr[largest]) {
        largest = left;
    }
    if (right < heap->size && heap->arr[right] > heap->arr[largest]) {
        largest = right;
    }
    if (largest != i) {
        int tmp = heap->arr[i];
        heap->arr[i] = heap->arr[largest];
        heap->arr[largest] = tmp;
        heapify(heap, largest);
    }
}

int *topk(int *nums, int n, int k) {
    heap_t *heap = heap_create(k);
    for (int i = 0; i < n; i++) {
        heap_insert(heap, nums[i]);
    }
    int *res = malloc(sizeof(int) * k);
    for (int i = 0; i < k; i++) {
        res[i] = heap_pop(heap);
    }
    heap_free(heap);
    return res;
}

int main() {
    int nums[] = {1, 3, 5, 2, 4, 6, 7, 9, 8};
    int n = sizeof(nums) / sizeof(nums[0]);
    int k = 3;
    int *res = topk(nums, n, k);
    for (int i = 0; i < k; i++) {
        printf("%d ", res[i]);
    }
    free(res);
    return 0;
}

Js实现：

class Heap {
    constructor(compare = (a, b) => a - b) {
        this.heap = [];
        this.compare = compare;
    }

    insert(val) {
        this.heap.push(val);
        this.heapifyUp(this.heap.length - 1);
    }

    pop() {
        if (this.heap.length === 0) {
            return null;
        }
        const root = this.heap[0];
        this.heap[0] = this.heap[this.heap.length - 1];
        this.heap.pop();
        this.heapifyDown(0);
        return root;
    }

    peek() {
        return this.heap[0] || null;
    }

    heapifyUp(i) {
        while (i > 0) {
            const parent = Math.floor((i - 1) / 2);
            if (this.compare(this.heap[i], this.heap[parent]) > 0) {
                [this.heap[i], this.heap[parent]] = [this.heap[parent], this.heap[i]];
            }
            i = parent;
        }
    }

    heapifyDown(i) {
        while (i < this.heap.length) {
            const left = 2 * i + 1;
            const right = 2 * i + 2;
            let largest = i;
            if (left < this.heap.length && this.compare(this.heap[left], this.heap[largest]) > 0) {
                largest = left;
            }
            if (right < this.heap.length && this.compare(this.heap[right], this.heap[largest]) > 0) {
                largest = right;
            }
            if (largest === i) {
                break;
            }
            [this.heap[i], this.heap[largest]] = [this.heap[largest], this.heap[i]];
            i = largest;
        }
    }

    topk(nums, k) {
        const heap = new Heap((a, b) => b - a);
        for (const num of nums) {
            heap.insert(num);
            if (heap.heap.length > k) {
                heap.pop();
            }
        }
        return heap.heap;
    }
}

const nums = [1, 3, 5, 2, 4, 6, 7, 9, 8];
const k = 3;
const res = new Heap().topk(nums, k);
console.log(res);

Rust实现：

use std::collections::BinaryHeap;

fn topk(nums: &[i32], k: usize) -> Vec<i32> {
    let mut heap = BinaryHeap::new();
    for num in nums {
        heap.push(*num);
        if heap.len() > k {
            heap.pop();
        }
    }
    heap.into_sorted_vec()
}

fn main() {