堆——实现及应用

Posted on 2020-09-17 In Coding Symbols count in article: 2.7k Reading time ≈ 2 mins.

概念

完全二叉树：若设二叉树的深度为 h，除第 h 层外，其它各层（1～h-1）的结点数都达到最大个数，第 h 层所有的结点都连续集中在最左边。
用数组表示（下标从 1 开始），则有：
- arr[i] 的左孩子是 arr[2*i]，右孩子是 arr[2*i+1]；
- arr[i] 的父节点是 arr[i/2]。
大顶堆：每个结点的值都大于或等于其左右孩子结点的值。
小顶堆：每个结点的值都小于或等于其左右孩子结点的值。

操作

注：本节以小顶堆为例，记堆的大小为 n。

首先定义一个堆。

class Heap {
private:
    int arr[maxn];
    int n;
    
    void shift_up(int i);
    void shift_down(int i);

public:
    Heap() {
        memset(arr, 0, sizeof(int) * maxn);
        n = 0;
    }

    void push(int x);
    void pop();
    int top();
    int size();
    bool empty();
};

上浮

从当前结点开始，和它的父节点比较：

若比父节点小则交换，然后将当前节点下标更新为原父节点下标；
否则退出。

void shift_up(int i) {
    while (i > 1 && arr[i] < arr[i>>1]) {
        swap(arr[i], arr[i>>1]);
        i >>= 1;
    }
}

下沉

当前节点与其左右孩子（如果有的话）中较小者作比较：

若后者比父节点小则交换，并更新当前节点下标为被交换的孩子节点下标；
否则退出。

void shift_down(int i) {
    while ((i << 1) <= n) {
        int j = i << 1;
        if (j < n && arr[j+1] < arr[j]) j++;
        if (arr[i] > arr[j]) swap(arr[i], arr[j]);
        else break;
        i = j;
    }
}

插入

向数组末尾插入新节点，然后使它上浮。

void push(int x) {
    arr[++n] = x;
    shift_up(n);
}

弹出

用尾节点覆盖根节点，堆大小减一，然后让新的根节点下沉。

void pop() {
    arr[1] = arr[n--];
    shift_down(1);
}

取顶

返回数组第一个元素。

1
2
3

int top() {
    return arr[1];
}

应用

堆排序

堆排序的基本思想是：将待排序序列构造成一个大顶堆，此时整个序列的最大值就是堆顶的根节点。将其与末尾元素进行交换，此时末尾就为最大值。然后将剩余 n-1 个元素重新构造成一个堆，这样会得到 n 个元素的次小值。如此反复执行，便能得到一个有序序列。

显然，大顶堆得到升序序列，小顶堆得到降序序列。

算法步骤为：

构造初始堆。从最后一个非叶子结点 arr[n/2] 开始，自下而上进行下沉操作；
将堆顶元素与末尾元素交换，此时的末尾元素从堆中排除，然后再次下沉根节点；
反复执行步骤 2，直到整个序列有序。

void shift_down(int* arr, int i, int n) {
    while ((i << 1) <= n) {
        int j = i << 1;
        if (j < n && arr[j+1] > arr[j]) j++;
        if (arr[i] < arr[j]) swap(arr[i], arr[j]);
        else break;
        i = j;
    }
}

void heap_sort(int* arr, int n) {
    // init heap
    for (int i = n >> 1; i >= 1; i--)
        shift_down(arr, i, n);
    // shift down from bottom to top
    while (--n) {
        swap(arr[1], arr[n+1]);
        shift_down(arr, 1, n);
    }
}

堆排序是一种选择排序，整体主要由构建初始堆+交换堆顶元素和末尾元素并重建堆两部分组成。其中构建初始堆经推导复杂度为 \(O(n)\)，在交换并重建堆的过程中，需交换 \(n-1\) 次，而重建堆的过程中，根据完全二叉树的性质， \([log2(n-1),log2(n-2)...1]\) 逐步递减，近似为 \(nlogn\) 。所以堆排序时间复杂度一般认为就是 \(O(nlogn)\) 级。

最小/大的 K 个数

用一个大根堆实时维护数组的前 \(k\) 小值。首先将前 \(k\) 个数插入大根堆中，随后从第 \(k+1\) 个数开始遍历，如果当前遍历到的数比大根堆的堆顶的数要小，就把堆顶的数弹出，再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。

反之，利用小根堆可以得到最大的 k 个数。

C++ 中的优先队列本质上就是由堆实现的，且默认是大根堆。而 Python 中的堆为小根堆，因此我们要对数组中所有的数取其相反数，才能使用小根堆维护前 \(k\) 小值。

合并 K 个有序链表

这个问题如果直接对所有链表一起排序，复杂度为 \(O(NlogN)\)，其中 \(N\) 为 \(K\) 个链表所有元素的总数。

然而我们应该充分利用链表本身是有序的条件，并通过堆来解决这个问题，其步骤是：

把每个链表第一个元素插入到最小堆；
从堆中取出最小的元素添加到结果列表中；
再从拿出去的元素所在的那个链表中取出下一个元素放到堆中；
重复第 2 步跟第 3 步，我们可以保证所有元素添加到了结果列表中且有序。

这种解法的时间复杂度可以达到 \(O(NlogK)\)，而空间复杂度为 \(O(K)\)。

概念

操作

上浮

下沉

插入

弹出

取顶

应用

堆排序

最小/大的 K 个数

合并 K 个有序链表

参考资料