剑指 Offer II 076. 数组中的第 k 大的数字 （用堆来解决topK问题）

剑指 Offer II 076. 数组中的第 k 大的数字 （用堆来解决topK问题）

题目描述

解题思路

这道题是一个很基础也很经典的 topK 问题。借此机会，顺便整理一下堆排序的写法。

首先是调整堆的结构以及建堆的代码，以大顶堆为例：

//调整堆的结构使其满足根节点大于等于其孩子节点
void maxHeapify(vector<int>& a, int i, int heapSize) {
    int left = i * 2 + 1, right = i * 2 + 2;	//完全二叉树下标为i的节点的左右孩子节点下标
    int largest = i;	//假定根结点和左右孩子中最大的是根
    //寻找根节点和左右孩子中的最大值
    if (left < heapSize && a[left] > a[largest]) {
        largest = left;
    }
    if (right < heapSize && a[right] > a[largest]) {
        largest = right;
    }
    // 如果跟节点不是最大值，调整堆的结构
    if (largest != i) {
        swap(a[i], a[largest]);	//使最大值成为跟节点
        maxHeapify(a, largest, heapSize);	//继续向下递归调整
    }
}
//从最后一个非叶子节点开始建堆
void buildMaxHeap(vector<int>& a) {
    //从完全二叉树最后一个非叶节点开始建堆
    for (int i = a.size() / 2 - 1; i >= 0; i--) {	
        maxHeapify(a, i, a.size());
    }
    //堆排序：建堆完成后，还需要从数组的最后倒着遍历，继续调整堆的结构
    for (int i = a.size() - 1; i >= 1; i--) {
        swap(a[i], a[0]);
        maxHeapify(a, 0, i);
    }
}

需要解释一下的是buildMaxHeap()函数。当建堆完成后，堆中的元素满足“根节点大于等于其孩子节点”这个原则，但是，每一层的元素之间不一定有序，也就是说这时候的数组中的元素还是部分有序的状态，如果要想得到有序数组则需要继续进行调整。如上面的代码中所示，从数组的末尾n - 1倒着遍历数组，将末尾元素a[n - 1]与堆顶元素a[0]交换，然后调整a[0]到a[n - 2]范围内堆的结构（末尾的元素已经有序，不需要调整），这样数组末尾的元素就是最大的。以此类推，最终得到一个非递减有序数组。

而针对此类 topK 问题，如果建堆的大小为 k ，那么堆内元素无须完全有序，只需要保证堆顶元素是这 k 个元素中最大（或最小）的即可，因此不需要完全调整至有序状态。也就是此题的解法，使用大小为 k 的小顶堆，代码如下：

//调整堆的结构使其满足根节点小于等于其孩子节点
void minHeapify(vector<int>& a, int i, int heapSize) {
    int left = i * 2 + 1, right = i * 2 + 2;
    int minIndex = i;
    if (left < heapSize && a[left] < a[minIndex]) {
        minIndex = left;
    }
    if (right < heapSize && a[right] < a[minIndex]) {
        minIndex = right;
    }
    if (minIndex != i) {
        swap(a[minIndex], a[i]);
        minHeapify(a, minIndex, heapSize);
    }
}
//从最后一个非叶子节点开始建堆，注意堆的大小不再是数组a的大小，而是定义的heapSize
void buildMinHeap(vector<int>& a, int heapSize) {
    //先用数组中前heapSize个数建堆
    for (int i = heapSize / 2 - 1; i >= 0; i--) {
        minHeapify(a, i, heapSize);
    }
    //再遍历数组中剩余的元素，如果它大于等于堆顶元素，就与堆顶元素交换，调整堆的结构
    for (int i = heapSize; i <= 1; i++) {
        if (a[i] >= a[0]) {
            swap(a[i], a[0]);
            minHeapify(a, 0, heapSize);
        }
    }
    //遍历完成后，堆内元素就是前k个最大的元素，堆顶元素就是第k大的元素
}

int findKthLargest(vector<int>& nums, int k) {
    buildMinHeap(nums, k);
    return nums[0];
}