Skip to content

数组中的第 K 个最大元素 ​

难度:⭐⭐⭐ 面试超高频 ​

考点 ​

  • 快速选择算法(QuickSelect)
  • partition 思想的变体应用
  • 堆方法作为备选
  • 字节面试必考题

题目描述 ​

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

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

函数签名 ​

go
func findKthLargest(nums []int, k int) int

示例 ​

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

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

要求 ​

  1. 平均时间复杂度 O(n)
  2. 使用快速选择算法(QuickSelect)
  3. 可接受修改原数组

提示 ​

  • 快速选择 = 快排只递归一半
  • 第 k 大 → 排序后的索引为 len(nums) - k
  • partition 后:pivot 在最终位置 i
    • i == target → 找到了
    • i < target → 在右半部分递归
    • i > target → 在左半部分递归
  • 随机 pivot 保证平均 O(n)

参考答案(Go) ​

点击展开参考答案
go
//go:build ignore

package answer

import "math/rand"

func findKthLargest(nums []int, k int) int {
	target := len(nums) - k // 第k大 = 排序后索引 len-k
	return quickSelect(nums, 0, len(nums)-1, target)
}

func quickSelect(nums []int, left, right, target int) int {
	if left == right {
		return nums[left]
	}
	pivotIdx := left + rand.Intn(right-left+1)
	nums[pivotIdx], nums[right] = nums[right], nums[pivotIdx]

	pivot := nums[right]
	i := left
	for j := left; j < right; j++ {
		if nums[j] < pivot {
			nums[i], nums[j] = nums[j], nums[i]
			i++
		}
	}
	nums[i], nums[right] = nums[right], nums[i]

	if i == target {
		return nums[i]
	} else if i < target {
		return quickSelect(nums, i+1, right, target)
	} else {
		return quickSelect(nums, left, i-1, target)
	}
}

持续学习,持续构建。