【剑指offer】40. 最小的k个数

题目描述

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：
输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：
输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：
0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000

解题思路

1. 最大堆

设置一个最大堆，存储当前最小的k个数。每次有新数字时，如果最大堆内数字数目小于k，则直接将新数字插入堆中；否则，将新数字与最大堆堆顶的值（即最小的k个数中的最大值）进行比较，如果小于最大堆堆顶值，则将最大堆堆顶值弹出，将新数字加入堆中。直到所有数字遍历完毕，最大堆中的k个数字即为数组中最小的k个数。

最大堆方法的时间复杂度O(nlogk)，空间复杂度O(k)。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = {};
        if(arr==null || arr.length<=0 || k<=0)
            return res;
        
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(10, new Comparator<Integer>(){
            public int compare(Integer o1, Integer o2){
                return o2-o1;
            }
        });

        for(int i=0; i<arr.length; i++){
            if(maxHeap.size()<k){
                maxHeap.add(arr[i]);
            }
            else{
                int temp = maxHeap.peek();
                if(temp>arr[i]){
                    maxHeap.poll();
                    maxHeap.add(arr[i]);
                }
            }
        }

        int i=0;
        res = new int[maxHeap.size()];
        for(Integer in : maxHeap){
            res[i++] = in;
        }
        return res;
    }
}

2. 快速排序（最高效）

快排在这道题中不用把整个数组进行排序，只需切分得到第k个数，则其左侧的k-1个数都小于第k个数，又侧的（n-k）个数都大于k，因此可以直接返回前k个数。
快速排序方法的时间复杂度O(n)（如果是整个数组都排序则是O(nlogn)），空间复杂度O(1)。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(arr==null || arr.length<=k)
            return arr;

        quicksort(arr, 0, arr.length-1, k);
        int[] res = new int[k];
        for(int i=0; i<k; i++)
            res[i] = arr[i];
        return res;
    }

    public void quicksort(int[] arr, int start, int end, int k){
        int pivot_index = partition(arr, start, end);
        if(pivot_index==k)
            return;
        else if(pivot_index<k){
            quicksort(arr, pivot_index+1, end, k);
        }
        else{
            quicksort(arr, start, pivot_index-1, k);
        }
    }

    public int partition(int[] arr, int left, int right){
        int pivot = arr[left];

        while(left<right){
            while(left<right && arr[right]>pivot){
                right--;
            }
            if(left<right){
                arr[left++] = arr[right];
            }
            while(left<right && arr[left]<=pivot){
                left++;
            } 
            if(left<right){
                arr[right--] = arr[left];
            }
            arr[left] = pivot;
        }
        return left;
    }
}

3. 冒泡排序

冒泡排序每轮将最小的移到前面，执行k轮即得到最小的k个数。
冒泡排序算是这三个方法里面最慢的，时间复杂度O(k(n-k))。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int len = arr.length;

        for(int i=0; i<k; i++){
            for(int j=len-1; j>i; j--){		//从后往前
                if(arr[j]<arr[j-1]){
                    int temp = arr[j];
                    arr[j] = arr[j-1];
                    arr[j-1] = temp;
                }
            }
        }

        int[] res = new int[k];
        int i=0;
        while(i<k){	//返回排序后数组的前k个元素
            res[i] = arr[i];
            i++;
        }

        return res;
    }
}