剑指 Offer 40. 最小的 k 个数

easy 快排

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入 4、5、1、6、2、7、3、8 这 8 个数字，则最小的 4 个数字是 1、2、3、4。

示例 1：

1
2
输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

1
2
输入：arr = [0,1,2,1], k = 1
输出：[0]

排序

排序后取前 k 个数，很笨！

1
2
3
4
func getLeastNumbers(arr []int, k int) []int {
	sort.Ints(arr)
	return arr[:k]
}

快排思想

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
// 快排思想
func getLeastNumbers(arr []int, k int) []int {
	var quickSort func(nums []int, l, r, x int)
	quickSort = func(nums []int, l, r, x int) {
		if l > r {
			return
		}
		i, j := l, r
		pivotPos := true
		for i < j {
			if nums[i] > nums[j] {
				nums[i], nums[j] = nums[j], nums[i]
				pivotPos = !pivotPos
			}
			if pivotPos {
				j--
			} else {
				i++
			}
		}
		num := i - l + 1 // 本次划分 可以保证划分正确的个数为num
		if x == num {
			// 正好
			return
		} else if x < num {
			quickSort(nums, l, i-1, x)
		} else {
			quickSort(nums, i+1, r, x-num)
		}
	}
	if k == 0 {
		return []int{}
	}
	quickSort(arr, 0, len(arr)-1, k)
	return arr[:k]
}

堆

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
func getLeastNumbers3(arr []int, k int) []int {
	// 大顶堆
	var heapify func(nums []int, root, end int)
	heapify = func(nums []int, root, end int) {
		// 大顶堆堆化，堆顶值小一直下沉
		for {
			// 左孩子节点索引
			child := root*2 + 1
			// 越界跳出
			if child > end {
				return
			}
			// 比较左右孩子，取大值，否则child不用++
			if child < end && nums[child] <= nums[child+1] {
				child++
			}
			// 如果父节点已经大于左右孩子大值，已堆化
			if nums[root] > nums[child] {
				return
			}
			// 孩子节点大值上冒
			nums[root], nums[child] = nums[child], nums[root]
			// 更新父节点到子节点，继续往下比较，不断下沉
			root = child
		}
	}
	end := len(arr) - 1
	// 从最后一个非叶子节点开始堆化
	for i := end / 2; i >= 0; i-- {
		heapify(arr, i, end)
	}
	// 依次弹出元素，然后再堆化，相当于依次把最大值放入尾部
	for i := end; i >= 0; i-- {
		arr[0], arr[i] = arr[i], arr[0]
		end--
		heapify(arr, 0, end)
	}
	return arr[:k]
}