快速选择算法Golang实现

类似求TopK问题中最常用的算法中，从时间复杂度最高到中等再到最优分别有不同的做法。在之前的学习中只学到了使用堆来优化TopK问题，但是这样的时间复杂度只能做到O（Nlogk）的大小，其中k是堆的大小。有一种更好的办法是基于快速排序的思想去优化的算法，叫做快速选择算法，它的时间复杂度能够做到O(N)的时间复杂度。这里的思路是：每次通过随机取得一个分区键，假设题目要求数组按照从大到小排序，那么通过将分区键移动到头部start，然后从头部的下一个元素开始遍历数组，遇到比分区键大的元素就交换到分区键后的已排序的下标的下一个位置，该指针假设就叫做index。最后遍历结束后将index的值与start的值交换，此时分区键就被移动到了index指针所指的位置，那么index左边的元素都是比分区键要大的，此时再通过对比index - start 与k的大小关系就可以判断下一次递归要从哪个区间开始，从而减少遍历的次数。

该算法的时间复杂度分析：该算法的最坏时间复杂度是每次递归都相当于重新遍历一次数组，那么最坏的时间复杂度是O(N)，但是通过随机算法的优化，使得每次取到的分区键都是均匀分布的，那么平均每次遍历的次数就近似看做一半，总的时间复杂度就可以看做：O(N) + O(N / 2) + ... 是一个等比数列求和，将N提取出来，就相当于O(NC)，其中C是常数，所以可以看做平均的时间复杂度是O(N)，该算法最典型的应用是：1. Leetcode215题：数组中的第K个最大元素。2. 变形题Leetcode347题：前K个高频元素。都可以使用快速选择算法完成。其中，215题官方的快速选择算法太过于复杂，懒得去看了，可以参考一下我这个写法，比较容易理解，具体代码如下：

func findKthLargest(nums []int, k int) int {
    // 快速选择算法
    return quickSelect(nums, 0, len(nums) - 1, k)
}

func quickSelect(arr []int, start, end, k int) int {
    rand.Seed(time.Now().Unix())
    picked := rand.Int() % (end - start + 1) + start
    // 交换到start
    arr[start], arr[picked] = arr[picked], arr[start]
    
    pivot := arr[start]
    index := start
    for i := start + 1; i <= end; i++ {
        if arr[i] >= pivot {
            arr[index + 1], arr[i] = arr[i], arr[index + 1]
            index++
        }
    }
    // 从大到小排序
    arr[index], arr[start] = arr[start], arr[index]
    // 这里k <= index - start是因为不包括pivot的值在里面
    if k <= index - start {
        return quickSelect(arr, start, index - 1, k)
    } else if k == index - start + 1 {
        return arr[index]
    } else {
        return quickSelect(arr, index + 1, end, k - (index - start + 1))
    }
}