解法一:最小堆
- 复杂度:O(NlogK) + O(K)
- 特别适合处理海量数据
- 应该使用大顶堆来维护最小堆,而不能直接创建一个小顶堆并设置一个大小,企图让小顶堆中的元素都是最小元素。
维护一个大小为 K 的最小堆过程如下:在添加一个元素之后,如果大顶堆的堆顶大小大于该元素,那么需要将大顶堆的堆顶元素与该元素交换,重新调整最大堆。
public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
ArrayList<Integer> ret = new ArrayList<>();
if(k==0 || k>input.length)
return ret;
int[] ks = new int[k];
for(int i=0; i<k; i++){
ks[i] = input[i];
}
buildMaxHeap(ks);
for(int i=k; i<input.length; i++){
if(input[i] < ks[0]){
ks[0] = input[i];
maxHeap(ks, ks.length, 0);
}
}
for(int val : ks)
ret.add(val);
return ret;
}
private void buildMaxHeap(int[] array){
if(array==null || array.length<=1)
return;
int half = (array.length-1)/2;
for(int i=half; i>=0; i--){
maxHeap(array, array.length, i);
}
}
private void maxHeap(int[] array, int heapSize, int index){
int left = index*2+1;
int right = index*2+2;
int largest = index;
if(left<heapSize && array[left]>array[largest])
largest = left;
if(right<heapSize && array[right]>array[largest])
largest = right;
if(largest != index){
int temp = array[largest];
array[largest] = array[index];
array[index] = temp;
maxHeap(array, array.length, largest);
}
}
这里可以使用Java的优先队列PriorityQueue来维护最大堆,代码如下:
public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
if(k==0 || k>input.length)
return new ArrayList<>();
PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1,o2)->o2-o1);
for(int val : input){
maxHeap.add(val);
if(maxHeap.size() > k){
maxHeap.poll();
}
}
return new ArrayList<>(maxHeap);
}
解法二:快速选择算法
- 复杂度:O(N) + O(1)
- 只有当允许修改数组元素时才可以使用
- 快速排序的 partition() 方法,会返回一个整数 j 使得 a[l…j-1] 小于等于 a[j],且 a[j+1…h] 大于等于 a[j],此时 a[j] 就是数组的第 j 大元素。可以利用这个特性找出数组的第 K 个元素,这种找第 K 个元素的算法称为快速选择算法。
public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
if(k==0 || k>input.length)
return new ArrayList<>();
int start = 0;
int end = input.length-1;
int index = partition(input, start, end);
while(index!=(k-1)){
if(index>k-1){
end = index-1;
index = partition(input, start, end);
}else{
start = index+1;
index = partition(input, start, end);
}
}
ArrayList<Integer> ret = new ArrayList<>();
for(int i=0; i<k; i++){
ret.add(input[i]);
}
return ret;
}
private int partition(int[] array, int begin, int end){
int low = begin;
int high = end;
int pivot = array[low];
while(low<high){
while(low<high && array[high]>=pivot)
high--;
array[low] = array[high];
while(low<high && array[low]<=pivot)
low++;
array[high] = array[low];
}
array[low] = pivot;
return low;
}