Medium | LeetCode 347. 前 K 个高频元素 | 快速排序

347. 前 K 个高频元素

给定一个非空的整数数组，返回其中出现频率前 *k* 高的元素。

示例 1:

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

示例 2:

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

提示：

• 你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。

• 你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。

• 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的。

• 你可以按任意顺序返回答案。

解题思路

方法一：优先队列

先使用HaspMap存放每个数字出现的次数, 把HashMap的每个键值对转为int[]数组, 以存放数值以及出现的次数, 然后使用优先队列可以很方便地取Top-K元素。

public int[] topKFrequent(int[] nums, int k) {
    // 先统计每个数字出现的次数
    Map<Integer, Integer> occurrences = new HashMap<Integer, Integer>();
    for (int num : nums) {
        occurrences.put(num, occurrences.getOrDefault(num, 0) + 1);
    }

    // int[] 的第一个元素代表数组的值，第二个元素代表了该值出现的次数, 并且建立的是小顶堆
    PriorityQueue<int[]> queue = new PriorityQueue<int[]>(new Comparator<int[]>() {
        public int compare(int[] m, int[] n) {
            return m[1] - n[1];
        }
    });
    for (Map.Entry<Integer, Integer> entry : occurrences.entrySet()) {
        int num = entry.getKey(), count = entry.getValue();
        if (queue.size() == k) {
            // 当堆大小为K时, 此时堆顶是K个元素最小值, 堆的其他元素均大于此堆顶值
            if (queue.peek()[1] < count) {
                // 如果入队的元素值大于堆顶元素, 则将此堆顶元素出队, 将更大的元素入队。
                queue.poll();
                queue.offer(new int[]{num, count});
            }
            // 如果入队的元素小于堆顶元素, 则直接丢弃, 因为比当前值大的至少有K个
        } else {
            // 队列大小小于K时直接入队
            queue.offer(new int[]{num, count});
        }
    }
    // 遍历完成后, 堆的所有元素就是Top-K元素
    int[] ret = new int[k];
    for (int i = 0; i < k; ++i) {
        ret[i] = queue.poll()[0];
    }
    return ret;
}

方法二： 快速排序

此题取Top-K的数和 Medium | LeetCode 215. 数组中的第K个最大元素 | 快速排序 基本相同。代码如下

public int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> count = new HashMap<>();
    for (int num : nums) {
        count.put(num, count.getOrDefault(num, 0) + 1);
    }
    List<int[]> freq = new ArrayList<>();
    for (Map.Entry<Integer, Integer> entry : count.entrySet()) {
        freq.add(new int[]{entry.getKey(), entry.getValue()});
    }
    int l = 0, r = freq.size() - 1;
    while (l < r) {
        int pos = partition(freq, l, r);
        if (pos < k-1) {
            r = pos + 1;
        } else if (pos > k - 1) {
            l = pos - 1;
        } else {
            break;
        }
    }
    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
        res[i] = freq.get(i)[0];
    }
    return res;
}

public int partition(List<int[]> nums, int left, int right) {
    int[] pivot = nums.get(right);
    int index = left;
    for (int i = index; i < right; i++) {
        if (nums.get(i)[1] >= pivot[1]) {
            Collections.swap(nums, i, index);
            index++;
        }
    }
    Collections.swap(nums, index, right);
    return index;
}