面试题40. 最小的k个数

一、使用快排解决TopK问题：O(N)

找前K大/前K小问题不需要对整个数组进行O(NlogN)的排序！
例如本题，直接通过快排切分排好第K小的数（下标为K-1），那么它左边的数就是比它小的另外K-1个数！

注意：在编写快排partition逻辑时，java需要额外编写一个swap函数，而不是像c++直接使用，并且swap函数中要传入数组的引用，否则最后把主元交换到数组某个位置是无法成功的！这也是java值传递的特点，这里花了一些时间，最后debug找出错误。

import java.util.Arrays;

class Solution{
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0)
            return new int[0];
        // k-1表示我们要找的是下标为k-1的数
        return quickSearch(arr, 0, arr.length - 1, k - 1);
    }

    private int[] quickSearch(int[] nums, int l, int r, int k) {
        // 每快排切分1次，找到排序后下标为j的元素，如果j恰好等于k就返回j以及j左边所有的数
        int j = partition(nums, l, r);
        if (j == k){
            return Arrays.copyOf(nums, j + 1);
        }
        // 否则根据下标j与k的大小关系来决定继续切分左段还是右段
        return j > k ? quickSearch(nums, l ,j - 1, k) : quickSearch(nums, j + 1, r, k);
    }

    // 快排切分，返回下标j，使得比nums[j]小的数都在j的左边，比nums[j]大的数都在j的右边
    private static int partition(int[] nums, int l, int r) {
        int v = nums[l];
        int j = l;
        for (int i = l + 1; i <= r; i++) {
            // 如果是需要求最大的k个数，只需要改为 nums[i] > k 即可
            if (nums[i] < v){
                swap(nums, j+1, i);
                j ++;
            }
        }
        swap(nums, l, j);
        return j;
    }

    // 自定义swap函数
    private static void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    public static void main(String[] args) {
        int[] temp = {1, 2, 4, 3, 5};
        int[] leastNumbers = new Solution().getLeastNumbers(temp, 3);
        for (int leastNumber : leastNumbers) {
            System.out.println(leastNumber);
        }
    }
}