【剑指Offer】最小的K个数（partition的应用和堆的应用）

法1：O(n) 要改变数组，partition的思想

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] ans = new int[k];
        findKMin(arr,0,arr.length-1,k);
        for(int i = 0; i < k; i++){
            ans[i] = arr[i];
        }
        return ans;
    }
    public void findKMin(int[] arr, int L, int R, int k){
        if(L < R){
            int[] p = partition(arr, L, R);
            if(p[0]+1==k){
                return;
            }
            findKMin(arr,L,p[0],k);
            findKMin(arr,p[1],R,k);
        }
    }
    public int[] partition(int[] arr, int L, int R){
        int less = L-1;
        int more = R+1;
        int basic = arr[R];
        while(L < more){
            if(arr[L] < basic){
                swap(arr,L++,++less);
            }else if(arr[L] > basic){
                swap(arr,L,--more);
            }else{
                L++;
            }
        }
        return new int[]{less,more};
    }
    public void swap(int[] arr, int i, int j){
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
}

 

法2：O(nlogK)+O(K)适合大数据量且不改变数组，使用堆（自己实现/使用PriorityQueue）

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>((o1,o2)->o2-o1);
        for(int i = 0; i < arr.length; i++){
            queue.add(arr[i]);
            if(queue.size() > k){
                queue.poll();
            }
        }
        int ans[] = new int[k];
        for(int i = 0; i < k; i++){
            ans[i] = queue.poll();
        }
        return ans;
    }
}