【每日算法Day 47】栈和队列2

今天主要来看优先队列（堆）相关的题目。一般k>1时涉及第k大或者前k个的问题可以考虑用堆，第k大用小顶堆第k小用大顶推，结合判断逻辑能将时间复杂度降到O(nlogk)。

LeetCode 215. 数组中的第K个最大元素 (opens new window)

题目描述

在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。

示例

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5

输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4

解题思路

看到题目可能最先想到使用任意一种效率较佳的排序算法排序后，返回倒数第K个位置的元素即可。

//使用API自带的快排
public int findKthLargest0(int[] nums, int k) {
    Arrays.sort(nums);
    return nums[nums.length-k];
}

//使用API自带的堆排序 大顶堆O(nlogn) 12ms
public int findKthLargest1(int[] nums, int k) {
    Queue<Integer> heap = new PriorityQueue<>((a,b)->b-a);
    for (int num : nums) {
        heap.offer(num);
    }
    int result=0;
    for (int i = 1; i <= k; i++) {
        result = heap.poll();
    }
    return result;
}

//小顶堆O(nlogk) 6ms
public int findKthLargest2(int[] nums, int k) {
    Queue<Integer> heap = new PriorityQueue<>();
    for (int num : nums) {
        if (heap.size() < k) {
            heap.offer(num);
        } else if (num > heap.peek()) {
            heap.poll();
            heap.offer(num);
        }
    }
    return heap.peek();
}

两种方法的时间复杂度都是O(nlogn)，空间复杂度如果不允许修改原数组的话都是O(n)。对于快排，其实有优化时间复杂度的空间。而对于这两种方法乃至于更多的排序算法，我们应该能够做到熟悉底层原理，在面试时自己写出实现（就算没有库函数写得那么好）。

LeetCode 347. 前 K 个高频元素 (opens new window)

题目描述

给定一个非空的整数数组，返回其中出现频率前 k 高的元素。你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。 你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的。 你可以按任意顺序返回答案。

示例

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

解题思路

首先自然是需要用哈希表统计出现频率。接下来需要进行排序，大多数常规的排序算法的时间复杂度最优也仅仅达到O(nlogn)，而题目要求时间复杂度优于O(nlogn)。所以需要尝试优化的算法。

法一，小顶堆：

可能你会想找频率最高的元素难道不应该是大顶堆吗？是呀，如果使用大顶堆就只能将数组中n个元素全部入堆才能确定前k个最大元素（因为对于一个元素e，你通过比较它和堆顶元素并不能确定这个元素是否为前k大的，只能将其入堆）。而使用小顶堆并限制堆的大小不超过k，则比较元素e和堆顶元素（前k个最大元素中最小的那个），如果小于堆顶元素则说明e一定不是前k个大元素就不用入堆维护了。这样因为堆的大小为k，时间复杂度可以降到O(nlogk)。

public int[] topKFrequent(int[] nums, int k) {
    //统计频次
    Map<Integer, Integer> map = new HashMap<>();
    for (int num : nums) {
        int count = map.getOrDefault(num, 0);
        map.put(num, count + 1);
    }

    //小顶堆排序
    Queue<Integer> heap = new PriorityQueue<>((a,b)->map.get(a)-map.get(b));
    for (Integer num : map.keySet()) {
        if (heap.size() < k) {
            heap.offer(num);
        } else if (map.get(num) > map.get(heap.peek())) {
            heap.poll();
            heap.offer(num);
        }
    }

    //输出结果
    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
        res[i]=heap.poll();
    }
    return res;
}

这样看昨天做的字节第二题好像用堆就多此一举了，因为要求找的只是统计出来的最大的那个元素，完全没必要放到堆中，直接用一个临时变量保存然后和Map.Entry逐个比较并更新就可以了。就算是要用堆，也是用小顶堆并限制大小为1，这样时间复杂度就是O(nlog1)即O(n)。

法二，桶排序：

public int[] topKFrequent2(int[] nums, int k) {
    //统计频次
    Map<Integer, Integer> map = new HashMap<>();
    for (int num : nums) {
        int count = map.getOrDefault(num, 0);
        map.put(num, count + 1);
    }

    //创建桶，以数字为桶的值，以频次为桶的下标，使用List以防有多个频次相同的数
    //一个数最多出现N次，则桶数组的长度为N+1（为了使用方便）
    List<Integer>[] buckets=new List[nums.length+1];
    map.forEach((key,value)->{
        if (buckets[value]==null)
            buckets[value]=new ArrayList<>();
        buckets[value].add(key);
    });

    //从后向前把所有非空的桶中所有的数往结果数组中填充，当结果数组已满时停止填充
    int[] res = new int[k];
    int index = 0;
    for (int i = nums.length; i >= 0; i--) {
        if (buckets[i] != null) {
            for (Integer num : buckets[i]) {
                if (index >= k)
                    break;
                res[index++] = num;
            }
        }
    }
    return res;
}

虽然填充的代码中有嵌套循环，不过那只是为了确保只填充k个数，该算法的时间复杂度为O(n)，空间复杂度为O(n)。