剑指offer面试题40：最小的k个数（golang实现）

题目描述

输入整数数组 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

算法分析

• 最容易想到的方法是直接排序，输出最小的k个数，用快排的话时间复杂度为O( n l o g n nlogn nlogn)。
• 也可以用快速排序的partition函数结合二分的思想来缩小区间，直到基准元素正好被放入数组中第k+1个位置，此时左边k个元素即为最小元素。

复杂度分析

问题规模为数组中元素的个数n

• 时间复杂度：O(n)， 最好时间复杂度为O( n n n)，即只需进行一趟排序；最坏时间复杂度为O( n n n)，每次递归后数组长度大致减半，求时间复杂度相当于等比数列求和，最坏时间复杂度约等于O( 2 n − 1 2n-1 2n−1)。所以均摊时间复杂度为O(n).
• 空间复杂度：O(logn)，平均递归深度为logn，每次递归有O(1)的空间消耗。

Golang代码如下

func partition(l, r, k int, arr []int) []int {
	pivot := arr[l]
	i, j := l, r
	for i < j {
		for i < j && arr[j] >= pivot {
			j--
		}
		arr[i] = arr[j]
		for i < j && arr[i] <= pivot {
			i++
		}
		arr[j] = arr[i]
	}
	arr[i] = pivot
	if i == k {
		return arr[:k]
	} else if i < k {
		return partition(i + 1, r, k, arr)
	} else {
		return partition(l, i - 1, k, arr)
	}
}

func getLeastNumbers(arr []int, k int) (res []int) {
	if k == len(arr) {
		return arr
	} else if k == 0 || len(arr) == 0 {
		return nil
	}
	return partition(0, len(arr) - 1, k, arr)
}