解
用PriorityQueue的话,极度简单,以前的几道题已经做过无数遍了,直接上答案。
public class Solution {
/*
* @param nums: an integer array
* @param k: An integer
* @return: the top k largest numbers in array
*/
public int[] topk(int[] nums, int k) {
int[] result = new int[k];
if (nums == null || nums.length < k) {
return result;
}
Queue<Integer> pq = new PriorityQueue<>();
for (int num : nums) {
pq.add(num);
if (pq.size() > k) {
pq.poll();
}
}
for (int i = k - 1; i >= 0; i--) {
result[i] = pq.poll();
}
return result;
}
}
解2
用QuickSelect做。比普通QuickSelect要多一个步骤,因为这道题不是要第K大的数了,而是要前K个最大的数,同时还要排序。如果不排序,QuickSelect已经做到了,因为在找到第K大的数的时候,已经保证了比它大的都在一边了,那么那一边加上它本身,就是前K个了。题目要求是排序的,所以要再对这一部分进行一下排序。
public class Solution {
/*
* @param nums: an integer array
* @param k: An integer
* @return: the top k largest numbers in array
*/
public int[] topk(int[] nums, int k) {
int[] result = new int[k];
if (nums == null || nums.length < k) {
return result;
}
quickSort(nums, 0, nums.length - 1, k);
// first k elements are the k largest elements
quickSort(nums, 0, k - 1, -1);
for (int i = 0; i < k; i++) {
result[i] = nums[i];
}
return result;
}
private void quickSort(int[] nums, int start, int end, int k) {
if (start >= end) {
return;
}
int left = start;
int right = end;
int pivot = nums[start + (end - start) / 2];
while (left <= right) {
while (left <= right && nums[left] > pivot) {
left++;
}
while (left <= right && nums[right] < pivot) {
right--;
}
if (left <= right) {
int temp = nums[left];
nums[left] = nums[right];
nums[right] = temp;
left++;
right--;
}
}
if (k != -1) {
if (start + k - 1 < right + 1) {
quickSort(nums, start, right, k);
} else if (start + k - 1 > right + 1) {
quickSort(nums, left, end, k - (left - start));
}
} else {
quickSort(nums, start, right, k);
quickSort(nums, left, end, k);
}
}
}
这里我为了省事,直接用QuickSort再sort了一遍,因为QuickSelect跟QuickSort的第一步都是一样的,我用了一个-1当作k的特殊值,如果是-1就用QuickSort,不是就QuickSelect。
分析
第一个方法,时间是O(nlogk),空间是O(k)。第二个方法,基于QuickSelect,空间常数O(1),时间第一步是O(n),再QuickSort前k个数字,即O(klogk),所以总共就是O(n+klogk)。这也很合理,如果k等于n,那么就等于sort了一遍array。