堆排序也是常见的一种排序算法,在生产中有很广泛的应用,比如优先级队列,TopK问题,生产中的TP99指标等。最近碰到了几个TopK问题,是如何用堆来解决的呢?比如:
堆是什么?
堆是一种完全二叉树,底层是用一维数组实现。分为大顶堆和小顶堆。以大顶堆为例,树中父节点不小于左右节点,并且以左右子节点为根的二叉树也都是大顶堆。比如下图a是大顶堆,b是小顶堆,c不是堆,因为不是完全二叉树。
堆的基本操作
堆的基本操作包括两个:插入节点insert和获取堆顶元素。
Insert
插入一般默认插入在数组最后,对应二叉树的话就是作为最后一个叶子节点。但插入后的二叉树可能不满足大顶堆或者小顶堆的性质了,需要将插入的节点放到堆中合适的位置,这时候需要将该节点层层上移(shiftUp),直到到合适的位置。
以大顶堆为例,如果插入的节点值比其父节点大,就交换二者位置,再继续考察父节点的位置是否需要上移。代码如下:
代码语言:javascript复制 void shiftUp(int k){
while (k > 1 && data[k / 2] < data[k]){
swap(data[k / 2], data[k]);
k /= 2;
}
}
getMax / getMin
以大顶堆为例,获取最大值很简单,直接把堆顶弹出即可。但是弹出后剩下的也必须要成为一个大顶堆。这时候把最后一个元素放到堆顶,但此时堆顶元素可能并不在其合适的位置,所以需要层层下移(shiftDown),直至移到合适的位置。
下移的时候首先需要考察该节点是否有左右孩子,如果有,那么再考察其值是否比孩子的值更小,如果更小,那么需要交换二者位置,继续往下考察;如果其值不必孩子的值更小,那说明该节点已经在合适的位置,此时堆已经是一个大顶堆了。
代码语言:javascript复制 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,代码如下:
代码语言:javascript复制for(int i=(arr.size()-2)/2;i>=0;i--){
shiftDown(arr, arr.size(), i);
}
如何解决TopK问题?
接下来回到本文最开始的问题,如何用堆来解决TopK问题?两步走!
- 构建堆:将原始数据构建成一个堆。
- 不断取堆顶:根据题目要求,取出堆顶。
面试题 17.14. 最小K个数
代码语言:javascript复制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小元素,与上一题差异在于这题的原始数据是用矩阵,即二维数组来表示。同样可以用堆来解决:
代码语言:javascript复制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];
}