搜索算法 - 数组中的第K个最大元素（快速选择算法）

介绍

若需要获取数组中的第K个最大元素，一般需要先排序，然后通过下标获取第K个最大元素。但其实可以基于快排的二分思路（即每次递归时都会将数组分为两个区间，且各区间都一定大于或小于中点），在排序选择中就可以确定第K个最大的元素。

代码参考自https://labuladong.github.io/algo/4/31/128/。

代码

func shuffle(arr []int) {
    var r int
    for i := len(arr) - 1; i >= 0; i-- {
        r = rand.Intn(i+1)
        arr[i], arr[r] = arr[r], arr[i]  
    }
}

func findKthLargest(nums []int, k int) int {
    var partition func(arr []int, left, right int)

    ans := -1

    partition = func(arr []int, left, right int) {
        p := arr[left]
        i := left
        j := right
        for {
            for i < j && arr[j] < p {
                j--
            }
            for i < j && arr[i] >= p {
                i++
            }
            if i >= j {
                break
            }
            arr[i], arr[j] = arr[j], arr[i]
        }
        arr[j], arr[left] = arr[left], arr[j]
      
        if j == k - 1 {
            ans = arr[j]
        } else if j > k - 1 {
            partition(arr, left, j-1)
        } else {
            partition(arr, j+1, right)
        }
    }

    // 打乱数组，防止时间复杂度恶化至n方
    shuffle(nums)
    partition(nums, 0, len(nums)-1)
    return ans
}

介绍

代码

相关题目