堆是如何解决TopK问题的?

2022-07-24 16:34:18 浏览数 (1)

堆排序也是常见的一种排序算法,在生产中有很广泛的应用,比如优先级队列,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问题?两步走!

  1. 构建堆:将原始数据构建成一个堆。
  2. 不断取堆顶:根据题目要求,取出堆顶。
面试题 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];
}

0 人点赞