如何求数组中最小的 k 个数？

问题

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

解法

排序再截取

最容易想到的方法就是先排序，再截取 k 个数，如：

// Go 1.19.3

func getLeastNumbers(arr []int, k int) []int {
	if len(arr) < k || k < 1 {
		return nil
	}
    // 对 arr 升序排序
	sort2.Ints(arr)
	return arr[:k]
}

上述算法中使用了标准库中的排序算法，在 Go 1.19.3 中采用了 pdqsort 这个新型算法（https://github.com/golang/go/issues/50154），比标准库中的原排序实现快2~60倍。

如果要求不能使用标准库排序的话，可以自己实现，较好的时间复杂度为 \(O(nlogn)\)，那么我们可以使用快速排序。

// 获取数组 arr 中 k 个最小的数
func getLeastNumbers(arr []int, k int) []int {
	if len(arr) < k || k < 1 {
		return nil
	}
	sort(&arr, 0, len(arr)-1)
	return arr[:k]
}

// 对数组 arr 中 [l,r] 区间的元素进行排序
func sort(arr *[]int, l, r int) {
	if l >= r {
		return
	}

	i := partition(arr, l, r)
	sort(arr, l, i-1)
	sort(arr, i+1, r)
}

// 获取数组 arr 中 [l,r] 区间查找基准元素，使比基准元素小的在基准元素左侧，比基准元素大的在基准元素右侧，返回基准元素位置
func partition(arr *[]int, l, r int) int {
	i, j := l, r
	for i < j {
		for i < j && (*arr)[j] >= (*arr)[l] {
			j--
		}
		for i < j && (*arr)[i] <= (*arr)[l] {
			i++
		}
		(*arr)[i], (*arr)[j] = (*arr)[j], (*arr)[i]
	}
	(*arr)[i], (*arr)[l] = (*arr)[l], (*arr)[i]
	return i
}

快速选择思路变形

由于题目要求找出数组中最小的 k 个数，并不要求我们进行完整排序，上面题解中使用 partition 函数时，最 k 小的数为基准数时其左侧都是比它小的，右侧都是比它大的，换句话说，如果找到的基准位置比 k-1小则最 k 小的在其右侧，反之在左侧，这样只要我们找到第 k 小的数的位置即可。

这种方法时间复杂度为 \(O(nlogn)\)。

func getLeastNumbers(arr []int, k int) []int {
	if len(arr) < k || k < 1 {
		return nil
	}
	// 存储结果
	var result = make([]int, k)
	// 区间起始索引
	l := 0
	// 区间结束索引
	r := len(arr) - 1
	pivot := partition(&arr, l, r)
	// 查找 pivot，使 pivot为第 k 小的数
	for pivot != k-1 && l < r {
		if pivot > k-1 {
			r--
		} else {
			l++
		}
		pivot = partition(&arr, l, r)
	}
	// 最终 pivot=k-1
	fmt.Println(pivot)
	// 最小的 k 个数在 [0,pivot] 区间内
	for i := 0; i <= pivot; i++ {
		result[i] = arr[i]
	}
	return result
}


// 获取数组 arr 中 [l,r] 区间查找基准元素，使比基准元素小的在基准元素左侧，比基准元素大的在基准元素右侧，返回基准元素位置
func partition(arr *[]int, l, r int) int {
	i, j := l, r
	for i < j {
		for i < j && (*arr)[j] >= (*arr)[l] {
			j--
		}
		for i < j && (*arr)[i] <= (*arr)[l] {
			i++
		}
		(*arr)[i], (*arr)[j] = (*arr)[j], (*arr)[i]
	}
	(*arr)[i], (*arr)[l] = (*arr)[l], (*arr)[i]
	return i
}

func main() {
	var l = []int{10, 15, 3, 6, 7, 12, 2, 1, 5}
	fmt.Println("original:", l)
    // result: [1 2 3 5]
	fmt.Println("result:", getLeastNumbers(l, 4))
}

这种方法修改了原数组，如果要求不能修改原数组的话，可以使用下面这个方法。

比较替换

设想我们有个容量为 k 的容器，用来表示最小的 k 个数。 先把数组前 k 个元素放到容器中作为最小的 k 个数； 遍历剩余的元素，如果元素小于容器中最大的数，用该元素替代容器中最大的数； 如果元素大于等于容器中最大的数，那么该数肯定不是最小的 k 个数中的一个，跳过；

这样遍历完成后，容器中就有最小的 k 个数，遍历数组的时间复杂度为 \(O(n)\)，从容器中找到和更新最大的数最优的时间复杂度是\(O(n)\)。

容器可以使用最大堆来实现。

Go 标准库中有最小堆的接口定义，如下：

// Go 1.19.3
// src/container/heap/heap.go

// The Interface type describes the requirements
// for a type using the routines in this package.
// Any type that implements it may be used as a
// min-heap with the following invariants (established after
// Init has been called or if the data is empty or sorted):
//
//	!h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

上面 sort.Interface 定义如下：

// Go 1.19.3
// src/sort/sort.go

// An implementation of Interface can be sorted by the routines in this package.
// The methods refer to elements of the underlying collection by integer index.
type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a transitive ordering:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a transitive ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

我们只要实现了 heap.Interface 就可以实现一个最小堆，要实现一个最大堆只要修改下 sort.Interface 中的 Less() 方法即可。

定义最大堆：

type IntMaxHeap struct {
	sort.IntSlice
}

/**
让 IntMaxHeap 实现heap.Interface 的方法
 */

// 重写sort的 Less 方法
func (h IntMaxHeap) Less(i, j int) bool {
	return h.IntSlice[i] > h.IntSlice[j]
}

// 实现堆的 Push 方法
func (h IntMaxHeap) Push(x interface{}) {
	h.IntSlice = append(h.IntSlice, x.(int))
}

// 实现堆的 Pop 方法
func (h IntMaxHeap) Pop() interface{} {
	x := h.IntSlice[len(h.IntSlice)-1]
	h.IntSlice = h.IntSlice[:len(h.IntSlice)-1]
	return x
}

实现求最小的 k 个数的方法：

func getLeastNumbers(arr []int, k int) []int {
	if len(arr) < k || k < 1 {
		return nil
	}
	// 最大堆的底层数组
	heapData := make([]int, k)
	// 给 heapData 填充数据
	copy(heapData, arr[:k])
	maxHeap := &IntMaxHeap{heapData}
	// 初始化 maxHeap，整理数据
	heap.Init(maxHeap)
	for i := k; i < len(arr); i++ {
		if x := arr[i]; x < maxHeap.IntSlice[0] {
            // 移除最大值
			heap.Pop(maxHeap)
            // 添加新的最大值
			heap.Push(maxHeap, x)
		}
	}
	return maxHeap.IntSlice
}

func main() {
	var l = []int{10, 15, 3, 6, 7, 12, 2, 1, 5}
	fmt.Println("original:", l)
	//  result: [5 3 1 2]
	fmt.Println("result:", getLeastNumbers(l, 4))
}

问题变型

如何找出数组中的中位数？

原问题的第一个题解中，本质思路是找到最 k 小的数，那么找到数组的中位数也可以用这个思路：

func findMedian(arr []int) int {
	if len(arr) < 1 {
		return -1
	}
	// 如果数组是有序的话，那么中位数的索引就是 len/2
	// 中位数索引
	midIdx := len(arr) / 2
	l := 0
	r := len(arr) - 1

	pivotIdx := partition(&arr, l, r)
	for pivotIdx != midIdx {
		if pivotIdx < midIdx {
			l++
		} else {
			r--
		}
		pivotIdx = partition(&arr, l, r)
	}
	return arr[pivotIdx]
}

func partition(arr *[]int, l, r int) int {
	i := l
	j := r
	for i < j {
		// 以 (*arr)[l] 作为基准
		for i < j && (*arr)[j] >= (*arr)[l] {
			// (*arr)[j] > (*arr)[l]，不用动，j 向左移
			j--
		}
		for i < j && (*arr)[i] <= (*arr)[l] {
			// (*arr)[i] < (*arr)[l]，不用动，i 向右移
			i++
		}
		// 这时 (*arr)[i] > (*arr)[j]，将两者值进行交换，使(*arr)[i] < (*arr)[j]
		(*arr)[i], (*arr)[j] = (*arr)[j], (*arr)[i]
	}
	// 交换 (*arr)[i], (*arr)[l]
	(*arr)[i], (*arr)[l] = (*arr)[l], (*arr)[i]

	return i
}

func main() {
	var l = []int{10, 15, 3, 6, 7, 12, 2, 1, 5}
	fmt.Println("original:", l)
	//  median: 6
	fmt.Println("median:", findMedian(l))
}

如何找出数组中出现次数超过数组长度一半的数？

如果数组是有序的话，将目标数作为一个窗口，那么窗口不管怎么左右滑动，窗口中数的值始终是数组的中位数，那么这个问题就和上面这个问题一样了，不过还要考虑不存在的情况。

func findTargetNumber(arr []int) int {
	if len(arr) < 1 {
		// 不存在
		return -1
	}
	var n int
	median := findMedian(arr)
	// 统计 median 出现的次数
	for i := 0; i < len(arr); i++ {
		if arr[i] == median {
			n++
		}
	}
	if n > len(arr)/2 {
		return median
	}
	// 不存在
	return -1
}