Top K 问题相信大家在面试过程中,经常被问到,下面就为大家来讲讲两种常见的实现算法。

一. 基于堆排序实现

• 1. 思路
       基于数组前K个数生成一个小顶堆，数组剩余元素依次与堆顶数据比较，小于等于堆顶数据时直接舍弃，大于堆顶数据时替换掉堆顶数据，并调整堆结构保证满足小顶堆要求。
• 2. 算法优势
       时间复杂度是O(N*logK)，不需要将数组一次全部加载到内存中，可以处理海量数据。

• 3. 图解

     
     
     Top K 小顶堆实现

• 4. 代码实现(Golang)

// 小顶堆法,找top k
func TopKByMinHeap(nums []int, k int) []int {
    length := len(nums)
    // 数组长度小于k,直接返回
    if length < k {     
        return nums
    }
        
    // 数组前k个数据取出,并生成小顶堆
    minHeap := make([]int, 0)
    minHeap = append(minHeap, nums[:k]...)
    CreateMinHeap(minHeap)

    // 遍历数组剩余数据,大于堆顶数据时,替换堆顶,重新维护小顶堆
    for i := k; i < length; i++ {
        if nums[i] > minHeap[0] {
            minHeap[0] = nums[i]
            MinHeapify(minHeap, 0, k)
        }
    }

    return minHeap
}

// 自底向上创建小顶堆
func CreateMinHeap(nums []int) {
    length := len(nums)
    for i := length - 1; i >= 0; i-- {
        MinHeapify(nums, i, length)
    }
}

// 维护小顶堆
func MinHeapify(nums []int, posIndex, length int) {
    // 堆左孩子节点index
    leftIndex := 2*posIndex + 1
    // 堆右孩子节点index
    rightIndex := 2*posIndex + 2
    // 当前节点以及左右孩子节点中最小值的index, 初始化为当前节点index
    minIndex := posIndex
    // 左孩子存在并且小于当前节点值时, 最小值index替换为左孩子index
    if leftIndex < length && nums[leftIndex] < nums[minIndex] {
        minIndex = leftIndex
    }
    // 右孩子存在并且小于当前节点值时, 最小值index替换为右孩子index
    if rightIndex < length && nums[rightIndex] < nums[minIndex] {
        minIndex = rightIndex
    }
    // 最小值节点index不等于当前节点时,替换当前节点和其中较小孩子节点值
    if minIndex != posIndex {
        nums[posIndex], nums[minIndex] = nums[minIndex], nums[posIndex]
        MinHeapify(nums, minIndex, length)
    }
}

二. 基于快排实现

• 1. 思路
       利用快排的分划函数找到分界点位置K，则前K个数据即所求结果。
• 2. 算法优势
       时间复杂度是O(N)，对于可以一次性加载到内存的数组效率很高。

• 3. 图解

     
     
     Top K 快排实现

• 4. 代码实现(Golang)

// 快排法,找Top k
func TopKByQuickSort(nums []int, k int) []int {
    length := len(nums)
    // 数组长度小于k,直接返回
    if length < k {
        return nums
    }

    // 数组进行快排,左侧边界
    left := 0
    // 数组进行快排,右侧边界
    right := length
    // 第一次快排后,获取分界点index
    pivotIndex := partition(nums, left, right)
    
    // 循环快排,找到分界点index刚好等于k
    for pivotIndex != k {
        if pivotIndex < k {
            // 分界点index小于k,继续对分界点右侧进行快排,重新获取分界点index
            left = pivotIndex + 1
        } else {
            // 分界点index大于k,缩小快排范围为上次分界点与本次分界点之间数据,重新获取分界点index
            right = pivotIndex
        }
        pivotIndex = partition(nums, left, right)
    }
    return nums[:k]
}

// 按分界点,进行快排,并返回分界点index
func partition(nums []int, left, right int) int {
    // 初始化分界值为左边界值
    pivot := nums[left]
    // 所有大于分界值的数据边界index
    pos := left
    
    // 小于分界值时,边界扩展,将数据替换到边界值index位置,
    for i := left; i < right; i++ {
        if nums[i] > pivot {
            pos++
            nums[i], nums[pos] = nums[pos], nums[i]
        }
    }
    
    // 交换分界值的数据边界index和分界点index,使得分界点左侧均大于分界点,右侧均小于分界点
    nums[pos], nums[left] = nums[left], nums[pos]

    return pos
}