数组 & 双指针 · 3/10

第 K 个最大元素

在无序数组中找到第 K 个最大的元素

问题

给定一个无序数组,找出其中第 K 个最大的元素。

例如:[3, 2, 1, 5, 6, 4],K = 2,返回 5。

解答

方法一:排序

最直观的方法,时间复杂度 O(n log n)。

function findKthLargest(nums, k) {
  // 降序排序后取第 k 个
  return nums.sort((a, b) => b - a)[k - 1];
}

// 测试
console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2)); // 5
console.log(findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4

方法二:快速选择

基于快排的分区思想,平均时间复杂度 O(n)。

function findKthLargest(nums, k) {
  // 第 k 大 = 第 n-k+1 小 = 索引 n-k
  const targetIndex = nums.length - k;
  return quickSelect(nums, 0, nums.length - 1, targetIndex);
}

function quickSelect(nums, left, right, targetIndex) {
  if (left === right) return nums[left];

  // 分区,返回基准元素的最终位置
  const pivotIndex = partition(nums, left, right);

  if (pivotIndex === targetIndex) {
    return nums[pivotIndex];
  } else if (pivotIndex < targetIndex) {
    // 目标在右半部分
    return quickSelect(nums, pivotIndex + 1, right, targetIndex);
  } else {
    // 目标在左半部分
    return quickSelect(nums, left, pivotIndex - 1, targetIndex);
  }
}

function partition(nums, left, right) {
  // 随机选择基准,避免最坏情况
  const randomIndex = left + Math.floor(Math.random() * (right - left + 1));
  swap(nums, randomIndex, right);

  const pivot = nums[right];
  let i = left;

  // 将小于 pivot 的元素放到左边
  for (let j = left; j < right; j++) {
    if (nums[j] < pivot) {
      swap(nums, i, j);
      i++;
    }
  }

  swap(nums, i, right);
  return i;
}

function swap(nums, i, j) {
  [nums[i], nums[j]] = [nums[j], nums[i]];
}

// 测试
console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2)); // 5

方法三:小顶堆

维护一个大小为 K 的小顶堆,时间复杂度 O(n log k)。

function findKthLargest(nums, k) {
  // 小顶堆,堆顶是最小值
  const heap = new MinHeap();

  for (const num of nums) {
    heap.push(num);
    // 保持堆大小为 k
    if (heap.size() > k) {
      heap.pop(); // 弹出最小的
    }
  }

  // 堆顶就是第 k 大
  return heap.peek();
}

class MinHeap {
  constructor() {
    this.data = [];
  }

  size() {
    return this.data.length;
  }

  peek() {
    return this.data[0];
  }

  push(val) {
    this.data.push(val);
    this._bubbleUp(this.data.length - 1);
  }

  pop() {
    const top = this.data[0];
    const last = this.data.pop();
    if (this.data.length > 0) {
      this.data[0] = last;
      this._bubbleDown(0);
    }
    return top;
  }

  _bubbleUp(index) {
    while (index > 0) {
      const parent = Math.floor((index - 1) / 2);
      if (this.data[parent] <= this.data[index]) break;
      [this.data[parent], this.data[index]] = [this.data[index], this.data[parent]];
      index = parent;
    }
  }

  _bubbleDown(index) {
    const length = this.data.length;
    while (true) {
      const left = 2 * index + 1;
      const right = 2 * index + 2;
      let smallest = index;

      if (left < length && this.data[left] < this.data[smallest]) {
        smallest = left;
      }
      if (right < length && this.data[right] < this.data[smallest]) {
        smallest = right;
      }
      if (smallest === index) break;

      [this.data[index], this.data[smallest]] = [this.data[smallest], this.data[index]];
      index = smallest;
    }
  }
}

// 测试
console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2)); // 5

关键点

  • 排序法简单但效率低,适合数据量小的场景
  • 快速选择平均 O(n),但最坏 O(n²),随机化可避免
  • 小顶堆稳定 O(n log k),适合 k 远小于 n 的场景
  • 第 K 大等价于第 (n - k + 1) 小,注意索引转换
  • 面试中优先写快速选择,展示算法功底