第 K 个最大元素

用堆排序如何求解第 K 个最大元素？

• 建大顶堆后, 每轮把堆顶换到堆底并缩小堆;
• 每轮交换都固定一个当前最大值, 第 k 轮交换后末尾已放置前 k 大的元素;
• 无需排序整个数组, 只需交换 k-1 次, 输出堆顶;
• 时间复杂度: 建堆 O(n) 加 k 次堆化 O(k log n);

function findKthLargest(nums: number[], k: number): number {
  const siftDown = (n: number, i: number) => {
    while (true) {
      const l = 2 * i + 1;
      const r = l + 1;
      let ma = i;
      if (l < n && nums[l] > nums[ma]) ma = l;
      if (r < n && nums[r] > nums[ma]) ma = r;
      if (ma === i) break;
      [nums[i], nums[ma]] = [nums[ma], nums[i]];
      i = ma;
    }
  };

  for (let i = nums.length - 1; i >= 0; i--) siftDown(nums.length, i);
  for (let i = nums.length - 1; i > nums.length - k; i--) {
    [nums[0], nums[i]] = [nums[i], nums[0]];
    siftDown(i, 0);
  }
  return nums[0];
}

用快速选择如何求解第 K 个最大元素？

• 每次哨兵划分都能确定基准数的最终下标 mid, 划分思路等同二分查找;
• 目标下标 target = n - k;
  • mid 等于 target: 直接返回 nums[mid];
  • mid 大于 target: 只在左半区间递归;
  • mid 小于 target: 只在右半区间递归;
• 每轮只处理一半, 平均 O(n), 最差 O(n^2);
• 划分前随机交换一个下标到 left 可避免有序输入退化, 见 快速排序;

function findKthLargestByPartition(nums: number[], k: number): number {
  const partition = (left: number, right: number): number => {
    const index = Math.floor(Math.random() * (right - left + 1) + left);
    [nums[left], nums[index]] = [nums[index], nums[left]];
    let i = left;
    let j = right;
    while (i < j) {
      while (i < j && nums[j] >= nums[left]) j--;
      while (i < j && nums[i] <= nums[left]) i++;
      [nums[i], nums[j]] = [nums[j], nums[i]];
    }
    [nums[i], nums[left]] = [nums[left], nums[i]];
    return i;
  };

  const select = (left: number, right: number): number => {
    if (left >= right) return nums[left];
    const mid = partition(left, right);
    const target = nums.length - k;
    if (mid === target) return nums[mid];
    return mid > target ? select(left, mid) : select(mid + 1, right);
  };

  return select(0, nums.length - 1);
}