【剑指offer】最小的K个数（Java解答，O(n)和O(nlogk)两种解答）

题目描述
输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。

O(n)算法，需要修改输入的数组

基本思路
通过随机快排算法，将数组随机分为两个部分，若分割点在第 k 个数的左边，则继续随机分割比分割点大的部分，反之，随机分割较小的那部分。循环以上步骤，直至分割点就是第k个数。

具体解答如下：

public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
    if (input == null || input.length < k || k == 0) {
        return new ArrayList<Integer>();
    }
    
    int start = 0;
    int end = input.length - 1;
    int index = partition(input, 0, input.length-1);
    while (index != k - 1) {
        if (index < k - 1) {
            start = index + 1;
        } else {
            end = index - 1;
        }
        index = partition(input, start, end);
    }
    
    ArrayList<Integer> result = new ArrayList<>(k);
    for (int i = 0; i < k; i++) {
        result.add(input[i]);
    }
    
    return result;
}

private int partition(int[] array, int start, int end) {
    if (array == null || array.length == 0 || start < 0 || end >= array.length) {
        throw new RuntimeException("Illegal arguments");
    }
    
    int index = new Random().nextInt(end-start+1) + start;
    int tmp = array[index];
    array[index] = array[start];
    array[start] = tmp;
    
    int smallPart = start;
    for (index = start+1; index <= end; index++) {
        if (array[index] < array[start]) {
            smallPart++;
            if (index != smallPart) {
                tmp = array[index];
                array[index] = array[smallPart];
                array[smallPart] = tmp;
            }
        }
    }
    
    tmp = array[start];
    array[start] = array[smallPart];
    array[smallPart] = tmp;
    
    return smallPart;
}

O(nlogk)的算法，适合处理海量数据

基本思路
借助一个大小为 k 的最大堆实现，轮询数组，若当前值比堆顶数字小，则清除堆顶元素，加入当前数；若比堆顶数字大，则继续轮询。

Java 中通过 PriorityQueue 实现最大堆，默认是最小堆，所以需要传入自定义的 Comparator。

具体解答如下：

public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
    if (input == null || k == 0 || input.length < k) {
        return new ArrayList<Integer>();
    }

    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, (i1, i2) -> i2 - i1);

    for (int value : input) {
        if (maxHeap.size() < k) {
            maxHeap.offer(value);
        } else {
            if (value < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(value);
            }
        }
    }

    return new ArrayList<Integer>(maxHeap);
}