堆是如何解决TopK问题的？
堆排序也是常见的一种排序算法，在生产中有很广泛的应用，比如优先级队列，TopK问题，生产中的TP99指标等。最近碰到了几个TopK问题，是如何用堆来解决的呢？比如：
堆是什么？
堆是一种完全二叉树，底层是用一维数组实现。分为大顶堆和小顶堆。以大顶堆为例，树中父节点不小于左右节点，并且以左右子节点为根的二叉树也都是大顶堆。比如下图a是大顶堆，b是小顶堆，c不是堆，因为不是完全二叉树。
堆的基本操作
堆的基本操作包括两个：插入节点insert和获取堆顶元素。
Insert
插入一般默认插入在数组最后，对应二叉树的话就是作为最后一个叶子节点。但插入后的二叉树可能不满足大顶堆或者小顶堆的性质了，需要将插入的节点放到堆中合适的位置，这时候需要将该节点层层上移（shiftUp），直到到合适的位置。
以大顶堆为例，如果插入的节点值比其父节点大，就交换二者位置，再继续考察父节点的位置是否需要上移。代码如下：
void shiftUp(int k){
while (k > 1 && data[k / 2] < data[k]){
swap(data[k / 2], data[k]);
k /= 2;
}
}getMax / getMin
以大顶堆为例，获取最大值很简单，直接把堆顶弹出即可。但是弹出后剩下的也必须要成为一个大顶堆。这时候把最后一个元素放到堆顶，但此时堆顶元素可能并不在其合适的位置，所以需要层层下移（shiftDown），直至移到合适的位置。
下移的时候首先需要考察该节点是否有左右孩子，如果有，那么再考察其值是否比孩子的值更小，如果更小，那么需要交换二者位置，继续往下考察；如果其值不必孩子的值更小，那说明该节点已经在合适的位置，此时堆已经是一个大顶堆了。
void shiftDown(int k){
// 当前孩子有孩子，就一定有左孩子（完全二叉树）
while (k * 2 <= count){
int j = k * 2;// 在此轮循环中，data[k]和data[j]交换位置
if (j + 1 <= count && data[j + 1] > data[j]){ // 当前孩子有右孩子
j = j + 1;
}
if (data[k] >= data[j]){
break; // k元素大于j元素的时候不必交换，终止循环
}
swap(data[k], data[j]);
k = j;
}
}建堆过程：heapify
给出一堆原始数据，如何构造成大顶堆呢？原始数据保存在一维数组里构成的一个完全二叉树不一定是一个堆，需要一步步调整各个节点的位置。
单独看二叉树的叶子节点，已经是一个大顶堆了。所以需要从这个完全二叉树的第一个非叶子节点开始，逐个向上调整各个非叶子节点的位置。构建堆的过程即heapify，代码如下：
for(int i=(arr.size()-2)/2;i>=0;i--){
shiftDown(arr, arr.size(), i);
}如何解决TopK问题？
接下来回到本文最开始的问题，如何用堆来解决TopK问题？两步走！
- 构建堆：将原始数据构建成一个堆。
- 不断取堆顶：根据题目要求，取出堆顶。
面试题 17.14. 最小K个数
void shiftDown(vector<int>&arr, int n, int k){
while (2 * k + 1<n){
int j = 2 * k + 1;
if (j + 1<n && arr[j + 1]<arr[j]){
j++;
}
if (arr[k] <= arr[j])break;
swap(arr[k], arr[j]);
k = j;
}
}
vector<int> smallestK(vector<int>& arr, int k) {
vector<int>ret;
if (k == 0)return ret;
// 建堆
for (int i = (arr.size() - 2) / 2; i >= 0; i--){
shiftDown(arr, arr.size(), i);
}
// 取k次堆顶
for (int i = arr.size() - 1; i >= 1; i--){
swap(arr[i], arr[0]);
ret.push_back(arr[i]);
k--;
if (k == 0)break;
shiftDown(arr, i, 0);
}
return ret;
}378. 有序矩阵中第K小的元素
这题是求有序矩阵中的第k小元素，与上一题差异在于这题的原始数据是用矩阵，即二维数组来表示。同样可以用堆来解决：
void __shiftDown(vector<vector<int>>& matrix, int n, int k){
int row = matrix.size();
while (2 * k + 1 < n){
int j = 2 * k + 1;
if (j + 1 < n && matrix[(j + 1) / row][(j + 1) % row] > matrix[j / row][j%row]){
j = j + 1;
}
if (matrix[k / row][k%row] >= matrix[j / row][j%row]){
break;
}
swap(matrix[k / row][k%row], matrix[j / row][j%row]);
k = j;
}
}
int kthSmallest(vector<vector<int>>& matrix, int k) {
int row = matrix.size();
int n = row*row;
// 建堆
for (int i = (n - 2) / 2; i >= 0; i--){
__shiftDown(matrix, n, i);
}
for (int i = n - 1; i>0; i--){
swap(matrix[i / row][i%row], matrix[0][0]);
__shiftDown(matrix, i, 0);
}
return matrix[(k - 1) / row][(k - 1) % row];
}腾讯云开发者
