一道面试题: Top K 问题

最近在招一个Go开发工程师，面试中时候我会问一个Top K的问题，这个问题是一个经典的面试题。

有时候我不会要求面试者写出答案，首先我听一下他的思想，如果写代码困难的话我都允许可以上网查标准库的文档，看看heap的用法。

相对来说比Redis的作者antirez的面试要轻松些了，他的面试题是要求面试者写出一个二叉搜索树。

这道题既然是经典题，很很多教科书或者算法网站上都有，比如leetcode也有，收录在Leetcode 算法题解精选一书中。

快速排序

我们可以采用快速排序的方法实现。
而且，因为我们只关心第K个最大的元素，我们只需要找出这个元素，不需要对整个数组进行完全排序。

// 从乱序的nums数组中找到第k大的数
func findKthLargest(nums []int, k int) int {
    n := len(nums)
    // 我们排序按照小的放在左边，大的放在右边的方式，所以第k大的数就是第n-k个数
    return quickselect(nums, 0, n - 1, n - k)
}
// 返回第k大的数
// l: 数组的左边界
// r: 数组的右边界
// k: 找第k大的数
func quickselect(nums []int, l, r, k int) int{
    if (l == r){ // =k
        return nums[k]
    }
    // 选取第一个数作为分区点，做一次调整
    partition := nums[l]
    i := l - 1
    j := r + 1
    for (i < j) {
        for i++;nums[i]<partition;i++{}
        for j--;nums[j]>partition;j--{}
        if (i < j) {
            nums[i],nums[j]=nums[j],nums[i]
        }
    }
    // 此时 i = j
    // 根据当前分区点的位置，判断k在左边还是右边
    // 如果k在分区点的左边，那么在左边继续找，右边界已经缩小到j了
    // 如果k在分区点的右边，那么在右边继续找，左边界已经缩小到j+1了
    // 通过递归，把边界逐步的缩小，直到左右两边界相等，也就是指向同一个位置，也就是函数开始的判断，就找到结果了
    if (k <= j){
        return quickselect(nums, l, j, k) // k 在左边, 则在左边继续找，右边界已经缩小到j了
    }else{ // k > j， k在右边
        return quickselect(nums, j + 1, r, k) // 从j+1到r找第k大的数,左边界已经缩小到j+1了
    }
}

递归一向都是难以理解，但是理解之后有觉得非常妙。

堆排序

这道题一个比较简单直观的解答就是堆排序。

我们使用一个最小堆，堆的大小为k，然后遍历数组，如果堆的大小小于k，就把元素放入堆中，如果堆的大小等于k，就把堆顶元素(当前堆内最小的元素)和当前元素比较，如果堆顶元素小于当前元素，就把堆顶元素弹出，把当前元素放入堆中。

// 从乱序的nums数组中找到第k大的数
func findKthLargest(nums []int, k int) int {
    n := len(nums)
    // 构建一个最小堆
    heap := make([]int, k)
    for i := 0; i < k; i++ {
        heap[i] = nums[i]
    }
    // 建堆
    for i := k/2 - 1; i >= 0; i-- {
        heapify(heap, i, k)
    }
    // 遍历数组，如果堆的大小小于k，就把元素放入堆中，如果堆的大小等于k，就把堆顶元素和当前元素比较，如果堆顶元素大于当前元素，就把堆顶元素弹出，把当前元素放入堆中。
    for i := k; i < n; i++ {
        if nums[i] > heap[0] {
            heap[0] = nums[i]
            heapify(heap, 0, k)
        }
    }
    return heap[0]
}
// 堆化
func heapify(heap []int, i, n int) {
    left := 2*i + 1
    right := 2*i + 2
    min := i
    if left < n && heap[left] < heap[min] {
        min = left
    }
    if right < n && heap[right] < heap[min] {
        min = right
    }
    if min != i {
        heap[i], heap[min] = heap[min], heap[i]
        heapify(heap, min, n)
    }
}

其实，Go标准库中提供了heap,我们自己不用实现。使用标准库中的heap,我们可以将上面的代码改写为：

import (
    "container/heap"
)
// 定义一个最小堆
type MinHeap []int
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}
func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}
func findKthLargest(nums []int, k int) int {
    h := &MinHeap{}
    heap.Init(h)
    for i := 0; i < k; i++ {
        heap.Push(h, nums[i])
    }
    for i := k; i < len(nums); i++ {
        if nums[i] > (*h)[0] {
            heap.Pop(h)
            heap.Push(h, nums[i])
        }
    }
    return (*h)[0]
}

虽然标准库实现了heap，但是实现的非常难用:

使用函数式的方式，没有面向对象的方式直观
要求Interface实现Len,Less,Swap三个方法(sort.Interface)，而且还要实现Push(x any)和Pop() any`两个方法。
包提供了heap.Init和heap.Fix、heap.Pop、heap.Push、heap.Remove等方法，Pop和Push和Interface的Pop和Push方法重名，这样会让人困惑。
而且，heap.Pop和Interface.Pop没有一点关系。heap.Push和Interface.Push没有一点关系。不要以为heap.Push直接调用Interface.Pop,虽然内部会调用，但是都有额外的处理