【劍指 Offer 題解】40. 最小的 K 個數

題目

輸入n個整數，找出其中最小的K個數。例如輸入4,5,1,6,2,7,3,8這8個數字，則最小的4個數字是1,2,3,4,。

思路

• 快速排序的 partition() 方法，會返回一個整數 j 使得 a[l…j-1] 小於等於 a[j]，且 a[j+1…h] 大於等於 a[j]，此時 a[j] 就是數組的第 j 大元素。可以利用這個特性找出數組的第 K 個元素，這種找第 K 個元素的算法稱爲快速選擇算法。
• 複雜度：O(N) + O(1)
• 只有當允許修改數組元素時纔可以使用

public ArrayList<Integer> GetLeastNumbers_Solution(int[] nums, int k) {
    ArrayList<Integer> ret = new ArrayList<>();
    if (k > nums.length || k <= 0)
        return ret;
    findKthSmallest(nums, k - 1);
    /* findKthSmallest 會改變數組，使得前 k 個數都是最小的 k 個數 */
    for (int i = 0; i < k; i++)
        ret.add(nums[i]);
    return ret;
}

public void findKthSmallest(int[] nums, int k) {
    int l = 0, h = nums.length - 1;
    while (l < h) {
        int j = partition(nums, l, h);
        if (j == k)
            break;
        if (j > k)
            h = j - 1;
        else
            l = j + 1;
    }
}

private int partition(int[] nums, int l, int h) {
    int p = nums[l];     /* 切分元素 */
    int i = l, j = h + 1;
    while (true) {
        while (i != h && nums[++i] < p) ;
        while (j != l && nums[--j] > p) ;
        if (i >= j)
            break;
        swap(nums, i, j);
    }
    swap(nums, l, j);
    return j;
}

private void swap(int[] nums, int i, int j) {
    int t = nums[i];
    nums[i] = nums[j];
    nums[j] = t;
}