当前位置：首页 > article >正文

【算法】用堆结构解决Top N问题

article 2024/12/5 4:09:00

在日常工作中，以及在算法练习和刷题时，Top N问题是经常遇到的。这里描述一种使用堆结构来处理Top N的问题。

二叉堆可以看成近似的完全二叉树。

堆的性质：
大根堆：除了根以外的所有节点i都要满足：A[PARENT(i)]≥A[i], 也就是说，某个节点至多与其父节点一样大。因此，堆中的最大元素存放在根节点中；并且，在任一子树中，该子树所包含的所有节点的值都不大于该子树根节点的值。
堆数据结构的最常用用法就是实现有点队列。
小根堆定义类似，将符号方向反过来（等号保留）即可。

第K近障碍物查询

在leetcode 3275 第K近障碍物查询中，就是一个典型的top N问题。
问题的描述：

有一个无限大的二维平面。
给你一个正整数 k ，同时给你一个二维数组 queries ，包含一系列查询：
queries[i] = [x, y] ：在平面上坐标 (x, y) 处建一个障碍物，数据保证之前的查询不会在这个坐标处建立任何障碍物。
每次查询后，你需要找到离原点第 k 近障碍物到原点的距离。
请你返回一个整数数组 results ，其中 results[i] 表示建立第 i 个障碍物以后，离原地第 k 近障碍物距离原点的距离。如果少于 k 个障碍物，results[i] == -1 。
注意，一开始没有任何障碍物。
坐标在 (x, y) 处的点距离原点的距离定义为 |x| + |y| 。

问题分析：从题目描述来看，待返回的结果要求queries的距离的第k近，也就是第k小。

常规列表排序

维护一个列表，用来存储至多k+1个元素，然后进行排序，取第k就是第k近的了。
代码实现如下：

class Solution {
    public int[] resultsArray(int[][] queries, int k) {
        int len = queries.length;
        int[] result  = new int[len];
        List<Integer> klist = new ArrayList<>();

        for(int i=0;i<len;i++){
           int distance = Math.abs(queries[i][0]) + Math.abs(queries[i][1]);
           klist.add(distance);
           if(klist.size()<k){
            result[i]=-1;
           }else{
            Collections.sort(klist);
            result[i]=klist.get(k-1);
            if(klist.size()>k){
                klist = klist.subList(0,k);
            }
           }
        }

        return result;
    }
}

这里的程序运行正常，但看起来运行效率不高。如果k很大，那么每次要排序的元素就会很多，耗时长。因为题目的约束里面，1 <= k <= 10^5

使用大根堆

在前面维护的是一个长度至多为k+1的列表，然后排序；但其实每次排序时，只会有一个元素的增量，这种正是堆结构的用武之地。
从需求分析来看，第k近，也就意味着如果取最小的k个元素构建堆，那么堆顶元素就是第k哥最小值，显然就是k近的元素。

class Solution {
    public int[] resultsArray(int[][] queries, int k) {
        int len = queries.length;
        int[] result = new int[len];
        PriorityQueue<Integer> pq = new PriorityQueue<>((o1, o2) -> o2.compareTo(o1));

        for (int i = 0; i < len; i++) {
            int distance = Math.abs(queries[i][0]) + Math.abs(queries[i][1]);
            pq.offer(distance);
            if (pq.size() < k) {
                result[i] = -1;
            } else {
                if (pq.size() > k) {
                    pq.poll();
                }
                result[i] = pq.peek();
            }

        }

        return result;
    }
}

效果对比

在使用列表排序时，碰到特殊（异常）case时，会极大的损耗性能。而使用堆结构时，由于堆的大小、每次操作的耗时都会小很多。

类似题目

前K个高频元素

leetcode 347. 前k个高频元素是这么定义的：

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按任意顺序返回答案。

问题分析：
这里可以看成是堆nums中的元素进行计数统计，然后求得top k值对应的计数值，然后计数值对应的元素即可。
这里是求前k高，采用小根堆。
在实现中，可以采用哈希表记录元素与计数，针对计数构建小根堆，在堆增加元素要大于k的时候，如果新元素小于等于堆顶元素则抛弃，如果新元素比堆顶元素大，则将堆顶元素出堆，并把新值入堆。

滑动窗口最大值

leetcode 239.滑动窗口最大值的题目描述如下：

给你一个整数数组 nums，有一个大小为 k 的滑动窗口从数组的最左侧移动到数组的最右侧。你只可以看到在滑动窗口内的 k 个数字。滑动窗口每次只向右移动一位。
返回滑动窗口中的最大值。

问题分析：
这里的要求是需要堆滑动窗口的k个值求最大值，特性是每次滑动一次，则会去掉一个值，而重新加入一个值。这就很符合堆结构的特征。
在具体实现时，继续构建一个大顶堆，但为了判断堆中的元素是否已经过了滑动窗口的界，还需要记录元素的索引值。
在窗口滑动时，将新元素和对应的索引入堆，这时候需要看堆顶元素是否在滑动窗口中，如果在，则堆顶元素就是最大值。
如果不在，需要持续将堆顶元素移除，直到找到的堆顶元素在在滑动窗口中。