215. 数组中的第K个最大元素

Problem

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。

请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

提示:

  • 1 <= k <= nums.length <= 105

  • -104 <= nums[i] <= 104

Solution

func findKthLargest(nums []int, k int) int {
    sink := func(heap []int, i int, end int) {
        for j := i * 2 + 1; j <= end; j = j * 2 + 1 {
            if j < end && heap[j+1] < heap[j] {
                j++
            }
            if heap[i] < heap[j] {
                break
            }
            heap[i], heap[j] = heap[j], heap[i]
            i = j
        }
    }

    //从下标 k-1 往前建立小顶堆
    for i := (k-1)>>1; i >=0; i-- {
        sink(nums, i, k-1)
    }

    //从下标k开始向堆中插入元素
    for i := k; i < len(nums); i++ {
        if nums[i] > nums[0] {
            nums[0], nums[i] = nums[i], nums[0]
            sink(nums, 0, k-1)
        }
    }

    return nums[0]
}

最后更新于

这有帮助吗?