剑指offer 40：最小的K个数

剑指offer 40：最小的K个数

基于partion函数的算法

这种方法会修改输入的数组。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k==0){
            return new int[0];
        }
        if(k==arr.length){
            return arr;
        }
        
        int len = arr.length;
        int l=0;
        int r = len-1;
        int mid = partion(arr,l,r);
        int[] res = new int[k];
        while(mid!=k){
            if(mid<k){
                l = mid+1;
            }else{
                r = mid-1;
            }
            mid = partion(arr,l,r);
        }
        
        for(int i=0;i<k;i++){
            res[i] = arr[i];
        }
        return res;

    }
    public int partion(int[] arr,int l,int r){
        int small = l-1;

        for(int i=l;i<r;i++){
            if(arr[i]<=arr[r]){
                small++;
                if(small!= i){
                    swap(arr,small,i);
                }
            }
        }

        small++;
        swap(arr,small,r);
        return small;

    }
    public void swap(int[] arr,int i,int j){
        if(i==j){
            return;
        }
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
}

时间复杂度$O(n)$。证明看《算法导论》。
空间复杂度$O(logn)$。递归调用。

堆

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        if(k==0){
            return res;
        }

        PriorityQueue<Integer> queque = new PriorityQueue(new Comparator<Integer>(){
            public int compare(Integer num1,Integer num2){
                return num2-num1;
            }
        });

        for(int i=0;i<k;i++){
            queque.add(arr[i]);
        }

        for(int i=k;i<arr.length;i++){
            if(arr[i]<queque.peek()){
                queque.poll();
                queque.add(arr[i]);
            }
        }

        for(int i=0;i<k;i++){
            res[i] = queque.poll();
        }
        return res;

    }
}

时间复杂度$O(nlogk)$，其中$n$是数组$arr$的长度。由于大根堆需要实时维护前$k$小值，所以插入删除都是$O(logk)$的时间复杂度，最坏情况下数组里$n$个数都会插入，所以共需要$O(nlogk)$的时间复杂度。
空间复杂度$O(k)$。因为大根堆里最多$k$个数。