当看到从一个集合中寻找最大/最小/最频繁的Top K个元素时,可以套用这个模式。
解决Top K问题的最好数据结构是堆(Heap)。这个模式也是使用大根堆/小根堆去解决Top K问题,流程如下:
- 插入K个元素到小根堆/大根堆中;
- 遍历剩下的元素,如果发现小于/或者大于堆中的元素,则将堆层的元素出队列,然后插入该元素。
- 最后留下的即为符合要求的元素。
从[3, 1, 5,12,2,11]中找到最大的3个元素的流程图如下:
如何识别符合这个模板的问题:
- 从一个给定集合中寻找最大/最小/最频繁的K个元素时;
- 被要求通过排序去找到一个要求的元素。
下面我们用模板解决问题。
https://leetcode.com/problems/kth-largest-element-in-an-array/
典型的top k问题,直接看代码吧:
public int findKthLargest(int[] nums, int k) {
PriorityQueue<Integer> pq = new PriorityQueue<>(); // 小根堆
for (int num : nums) {
if (pq.size() < k) {
pq.offer(num);
} else {
if (num > pq.peek()) {
pq.poll();
pq.offer(num);
}
}
}
return pq.poll();
}
时间和空间复杂度为:O(N lg K) running time + O(K) memory。
当然这道题我们也可以直接对数组就行排序,然后找到第K个元素:
public int findKthLargest(int[] nums, int k) {
Arrays.sort(nums);
return nums[nums.length - k];
}
时间和空间复杂度为:O(N lg N) running time + O(1) memory。
可以看到,当K远小于N时,第一种方法较好。
这道题还有更多的解决方法,具体的请看leetcode的评论:
https://leetcode.com/problems/kth-largest-element-in-an-array/discuss/60294/Solution-explained
https://leetcode.com/problems/k-closest-points-to-origin/
这道题是找到离原点(0, 0)的Top closest的点,明显符合我们的要求,使用Top k element模式。具体的看代码:
public int[][] kClosest(int[][] points, int k) {
int[][] res = new int[k][];
// 存储K closet point集合,大根堆
PriorityQueue<int[]> pq = new PriorityQueue<>(
(a, b) -> ((b[0] * b[0] + b[1] * b[1]) - (a[0] * a[0] + a[1] * a[1])));
// 先加K个元素到queue中,等元素超过K个则先入队列,然后将距离最远的出队列。最终队列中保存的都是符合要求的点
for (int[] point : points) {
if (pq.size() < k) {
pq.offer(point);
} else {
pq.offer(point);
pq.poll();
}
}
// 输出结果
int index = 0;
while (!pq.isEmpty()) {
res[index++] = pq.poll();
}
return res;
}
https://leetcode.com/problems/top-k-frequent-elements/
Top K最频繁,明显也符合该模式要求。唯一不同的地方在于我们需要使用一个hashmap去存储元素和它出现的频率,然后将频率作为筛选要素。
public int[] topKFrequent(int[] nums, int k) {
Map<Integer, Integer> frequency = new HashMap<>();
for (int n : nums) {
frequency.put(n, frequency.getOrDefault(n, 0) + 1);
}
PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparing(frequency::get)); // 使用map.get()作为排序,也就是频率
for (Integer i : frequency.keySet()) { // 入队列的为元素
pq.offer(i);
if (pq.size() > k) pq.poll();
}
// 返回结果
int[] result = new int[k];
while (k > 0) {
k--;
result[k] = pq.poll();
}
return result;
}