leetcode 面试题 17.14. 最小K个数

力扣题目链接1
力扣题目链接2
牛客题目链接

该题是一道经典题，有多种解法。上面题目之间略有不同，请注意看题目要求。以下代码都针对力扣题目链接1，全部提交通过。

解法一

从小到大整体排序，然后取前k个数。

import "sort"

func smallestK(arr []int, k int) []int {
    if k < 1 {
        return nil
    }
    l := len(arr)
    if l < k {
        return nil
    }
    if l == k {
        return arr
    }
    sort.Sort(sort.IntSlice(arr))
    return arr[:k]
}

假设基于快速排序，时间复杂度$O(nlogn)$，空间复杂度$O(logn)$。

解法二

基于小顶堆。先将数组原地调整为小顶堆，然后弹出堆顶的k个数。

func smallestK(arr []int, k int) []int {
    if k == 0 {
        return []int{}
    }
    l := len(arr)
    if l <= k {
        return arr
    }
    if l < 2 {
        return arr
    }
    e := l-1
    for i := l/2-1; i >= 0; i-- {
        heapify(arr, i, e)
    }
    var j int
    for i := l-1; i > 0; i-- {
        arr[i], arr[0] = arr[0], arr[i]
        j++
        if j == k {
            break
        }
        heapify(arr, 0, i-1)
    }
    return arr[l-k:]
}

// 小顶堆
func heapify(a []int, s, e int) {
    i := s
    temp := a[i]
    j := 2*i+1
    for i < e && j <= e {
        if j+1 <= e && a[j+1] < a[j] {
            j++
        }
        if a[j] >= temp {
            break
        }
        a[i] = a[j]
        i = j
        j = 2*i+1
    }
    a[i] = temp
}

时间复杂度$O(nlogn)$，空间复杂度$O(1)$。

解法三

基于大顶堆。大顶堆需要额外的存储空间，容纳k个数。顺序遍历数组，将元素压入堆，直到堆的有效数据个数达到k。然后继续遍历数组，后续每一个元素都和堆顶的数做比较，如果元素值比堆顶的数小就替换堆顶的数，并做堆调整，直到所有数组元素都遍历完为止。最后堆中那k个数就是结果。

func smallestK(arr []int, k int) []int {
    if k <= 0 {
        return nil
    }
    l := len(arr)
    if l < k {
        return nil
    }
    if l == k {
        return arr
    }
    // l > k
    h := NewMaxHeap(k)
    var i int
    for i < l {
        h.Push(arr[i])
        i++
        if h.Size() >= k {
            break
        }
    }
    for i < l {
        n := arr[i]
        if n < h.Top() {
            h.ReplaceTop(n)
        }
        i++
    }
    return h.Get()
}

// 大顶堆
type MaxHeap struct {
    data []int
    size int
}

func NewMaxHeap(k int) *MaxHeap {
    return &MaxHeap{
        data: make([]int, k),
    }
}

func (o *MaxHeap) Size() int {
    return o.size
}

func (o *MaxHeap) IsEmpty() bool {
    return o.size == 0
}

func (o *MaxHeap) Push(n int) {
    o.data[o.size] = n
    o.size++
    o.siftUp()
}

func (o *MaxHeap) siftUp() {
    if o.size < 2 {
        return
    }
    child := o.size-1
    childVal := o.data[child]
    for child > 0 {
        p := (child-1)/2 // parent
        pV := o.data[p]
        if pV >= childVal {
            break
        }
        o.data[child] = pV
        child = p
    }
    o.data[child] = childVal
}

func (o *MaxHeap) Top() int {
    if o.IsEmpty() {
        panic("MaxHeap is empty")
    }
    return o.data[0]
}

func (o *MaxHeap) ReplaceTop(n int) {
    if o.IsEmpty() {
        panic("MaxHeap is empty")
    }
    o.data[0] = n
    o.siftDown()
}

func (o *MaxHeap) siftDown() {
    if o.size < 2 {
        return
    }
    var p int
    pV := o.data[p]
    for {
        child := 2*p + 1
        if child >= o.size {
            break
        }
        if child+1 < o.size && o.data[child+1] > o.data[child] {
            child++
        }
        if pV >= o.data[child] {
            break
        }
        o.data[p] = o.data[child]
        p = child
    }
    o.data[p] = pV
}

func (o *MaxHeap) Get() []int {
    return o.data
}

时间复杂度$O(nlogk)$，空间复杂度$O(k)$。

解法四

快速选择算法。

• 递归实现

func smallestK(arr []int, k int) []int {
    if k <= 0 {
        return nil
    }
    l := len(arr)
    if k >= l {
        return arr
    }
    quickSelect(arr, 0, l-1, k)
    return arr[:k]
}

func quickSelect(arr []int, low, high, k int) {
    pivotIdx := doPivot(arr, low, high)
    if pivotIdx == k-1 {
        return
    }
    if pivotIdx < k-1 {
        low = pivotIdx+1
    } else if pivotIdx > k-1 {
        high = pivotIdx-1
    }
    quickSelect(arr, low, high, k)
}

func doPivot(arr []int, low, high int) int {
    medianOfThree(arr, low, high)
    pivot := arr[low]
    for low < high {
        for low < high && arr[high] >= pivot {
            high--
        }
        arr[low] = arr[high]
        for low < high && arr[low] <= pivot {
            low++
        }
        arr[high] = arr[low]
    }
    arr[low] = pivot
    return low
}

func medianOfThree(arr []int, low, high int) {
    mid := low+(high-low)>>1
    if arr[mid] > arr[high] {
        arr[mid], arr[high] = arr[high], arr[mid]
    }
    if arr[low] > arr[high] {
        arr[low], arr[high] = arr[high], arr[low]
    }
    if arr[low] < arr[mid] {
        arr[low], arr[mid] = arr[mid], arr[low]
    }
}

快速选择。递归。
时间复杂度：最优O(n)，最差O(n^2)。
空间复杂度：最优O(logn)，最差O(n)。

• 非递归实现

func smallestK(arr []int, k int) []int {
    if k == 0 {
        return []int{}
    }
    l := len(arr)
    if l < 2 {
        return arr
    }
    if l <= k {
        return arr
    }
    var (
        low int
        high = l-1
    )
    for low < high {
        pivotIndex := doPivot(arr, low, high)
        if pivotIndex+1 == k {
            break
        } else if pivotIndex+1 < k {
            low = pivotIndex+1 // 再在后一个区间找
        } else {
            high = pivotIndex-1 // 再在前一个区间找
        }
    }
    return arr[:k]
}

// 从小到大
func doPivot(a []int, low, high int) int {
    // 三数取中
    // high - low + 1 >= 3
    if high - low > 1 {
        m := low + (high-low)/2
        if a[m] > a[high] {
            a[m], a[high] = a[high], a[m]
        }
        if a[low] > a[high] {
            a[low], a[high] = a[high], a[low]
        }
        if a[m] > a[low] {
            a[m], a[low] = a[low], a[m]
        }
        // a[m] <= a[low] <= a[high]
    }
    pivot := a[low]
    for low < high {
        for low < high && a[high] >= pivot {
            high--
        }
        a[low] = a[high]
        for low < high && a[low] <= pivot {
            low++
        }
        a[high] = a[low]
    }
    a[low] = pivot
    return low
}

快速选择。非递归。
时间复杂度：最优O(n)，最差O(n^2)。
空间复杂度：因为就地处理，所以是O(1)。

解法五

BFPRT算法。

func smallestK(arr []int, k int) []int {
    if k < 1 {
        return nil
    }
    l := len(arr)
    if l == 0 {
        return nil
    }
    if l < k {
        return nil
    }
    if l == k {
        return arr
    }
    // l > k
    kthSmallest(arr, 0, l-1, k)
    return arr[:k]
}

func kthSmallest(nums[] int, low, high, k int) int {
    for low < high {
        pivotIndex := selectPivot(nums, low, high)
        pivotIndex = partition(nums, low, high, pivotIndex, k)
        if k-1 < pivotIndex {
            high = pivotIndex-1
        } else if k-1 == pivotIndex {
            return pivotIndex
        } else {
            low = pivotIndex+1
        }
    }
    return low
}

func selectPivot(nums []int, low, high int) int {
    if high-low+1 <= 5 {
        return partition5(nums, low, high)
    }
    // median of medians
    t := low
    for i := low; i <= high; i += 5 {
        j := i+4
        if j > high {
            j = high
        }
        m := partition5(nums, i, j)
        t = low + (i-low)/5
        nums[m], nums[t] = nums[t], nums[m]
    }
    mid := low + (t-low)>>1
    if (t-low+1)&1 == 0 {
        mid++ // 下中位数
    }
    return kthSmallest(nums, low, t, mid)
}

func partition(nums []int, low, high, pivotIndex, kth int) int {
    // a three-way partition
    var (
        i = low
        j = low
        k = high
        pivot = nums[pivotIndex]
    )
    for j <= k {
        n := nums[j]
        if n < pivot {
            nums[i], nums[j] = n, nums[i]
            i++
            j++
        } else if n > pivot {
            nums[j], nums[k] = nums[k], n
            k--
        } else {
            j++
        }
    }
    if kth-1 < i {
        return i
    } else if kth-1 < j {
        return kth-1
    } else {
        return k
    }
}

func partition5(nums []int, low, high int) int {
    // 从小到大，直接插入排序
    for i := low+1; i <= high; i++ {
        n := nums[i]
        j := i-1
        for j >= 0 && nums[j] > n {
            nums[j+1] = nums[j]
            j--
        }
        nums[j+1] = n
    }
    return low + (high-low)>>1
}

时间复杂度$O(n)$，但是n通常有一个较大的常数系数C。