leetcode-剑指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]

方法一：快排

解题思路

• 快排，找基准。然后从右边找一个比基准大的，从左边找一个比基准小的。交换。基准和小的交换。这样就将数组以基准分为两部分，递归处理这两部分（分治）
• 注意先从右边找，再从左边找。
• 排完序取前k个元素返回即可

代码实现

func getLeastNumbers(arr []int, k int) []int {
    quickSort(arr, 0, len(arr) - 1)
    return arr[:k]
}

func quickSort(arr []int, left, right int) {
    if left >= right {
        return
    }

    i := left
    j := right

    for i < j {
        for i < j && arr[j] >= arr[left] {
            j--
        }
        for i < j && arr[i] <= arr[left] {
            i++
        }
        arr[i], arr[j] = arr[j], arr[i]
    } 
    arr[i], arr[left] = arr[left], arr[i]

    quickSort(arr, left, i - 1)
    quickSort(arr, i + 1, right)
}

方法二：改进快排

解题思路

• 题目只要求返回最小的 k 个数，对这 k 个数的顺序并没有要求。因此，只需要将数组划分为 最小的 kk 个数 和 其他数字 两部分即可，而快速排序的哨兵划分可完成此目标。
• 根据快速排序原理，如果某次哨兵划分后 基准数正好是第 k+1k+1 小的数字 ，那么此时基准数左边的所有数字便是题目所求的 最小的 k 个数 。

代码实现

func getLeastNumbers(arr []int, k int) []int {
    if k >= len(arr) {
        return arr
    }
    return quickSort(arr, k, 0, len(arr) - 1)
}

func quickSort(arr []int, k, left, right int) []int {
    i := left
    j := right

    for i < j {
        for i < j && arr[j] >= arr[left] {
            j--
        }
        for j < j && arr[i] >= arr[left] {
            i++
        }
        arr[i], arr[j] = arr[j], arr[i]
    }
    arr[i], arr[left] = arr[left], arr[i]

    if i < k {
        return quickSort(arr, k, i + 1, right)
    }
    if i > k {
        return quickSort(arr, k, left, i - 1)
    }
    return arr[:k]
}