跳到主要内容

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

mid

给定整数数组 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

快排思想

详见 https://leetcode.cn/problems/kth-largest-element-in-an-array/solution/partition-zhu-shi-by-da-ma-yi-sheng/

func findKthLargest(nums []int, k int) int {
var partition func(nums []int, l, r int) int
partition = func(nums []int, l, r int) int {
pivotPos := true // 基准值位置 true为前
i, j := l, r
for i < j {
if nums[i] < nums[j] {
// 逆序
pivotPos = !pivotPos
nums[i], nums[j] = nums[j], nums[i]
}
if pivotPos {
j--
} else {
i++
}
}
return i
}

left := 0
right := len(nums) - 1

for left < right {
p := partition(nums, left, right)
if p+1 == k {
return nums[p]
} else if p+1 < k {
left = p + 1
} else {
right = p - 1
}
}
return nums[right]
}

partition 的另一种写法

partition = func(nums []int, l, r int) int {
pivot := nums[r]
for i := l; i < r; i++ {
if nums[i] > pivot {
nums[l], nums[i] = nums[i], nums[r]
l++
}
}
nums[l], nums[r] = nums[r], nums[l]
return l
}