【LeetCode 热题100道笔记】数组中的第K个最大元素

题目描述

给定整数数组 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<=10^5$
• $-10^4<=nums[i]<=10^4$

思考一

通过构建大顶堆，利用堆的特性（堆顶元素始终是最大值）来快速定位第k个最大元素。具体来说，我们先将所有元素加入大顶堆，然后弹出堆顶元素k-次，此时的堆顶元素就是第k个最大元素。

算法过程

1. 构建大顶堆：将所有元素插入大顶堆（堆顶始终为当前最大值）
2. 提取前k-1个最大值：执行k-1次弹出堆顶操作
3. 获取结果：此时堆顶元素即为第k个最大元素

时空复杂度

• 时间复杂度：O(n log n)

  • 插入n个元素：O(n log n)（单次插入O(log n)）
  • 弹出k-1个元素：O(k log n)（单次弹出O(log n)）

• 空间复杂度：O(n)

  • 需存储所有元素在堆中，占用O(n)空间

代码

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function(nums, k) {
    const priorityQueue = new MyPriorityQueue(nums.length);
    for (let num of nums) {
        priorityQueue.push(num);
    }

    for (let i = 0; i < k-1; i++) {
        priorityQueue.pop();
    }

    return priorityQueue.front();    
};

class MyPriorityQueue {

    constructor(capacity = 1000, compare = (a, b) => a - b) {
        this._data = [];
        this._capacity = capacity;
        this._size = 0;
        this._compare = compare;
    }

    front() {
        return this._data[0];
    }

    push(num) {
        if (this._capacity === this._size) {
            this.pop();
        }
        this._data.push(num);
        this.swim();
        this._size++;
    }

    pop() {
        if (this._data.length === 0) return;
        [this._data[0], this._data[this._data.length-1]] = [this._data[this._data.length-1], this._data[0]];
        const item = this._data.pop();
        this.sink();
        this._size--;
        return item;
    }

    swim(index = this._data.length-1) {
        while (index > 0) {
            let pIndex = Math.floor((index-1)/2);
            if (this._compare(this._data[index],this._data[pIndex]) > 0) {
                [this._data[index], this._data[pIndex]] = [this._data[pIndex], this._data[index]];
                index = pIndex;
                continue;
            }
            break;
        }

    }

    sink(index = 0) {
        const n = this._data.length;
        while (true) {
            let left = 2 * index + 1;
            let right = left + 1;
            let biggest = index;

            if (left < n && this._compare(this._data[left], this._data[index]) > 0) {
                biggest = left;
            }
            if (right < n && this._compare(this._data[right], this._data[biggest]) > 0) {
                biggest = right;
            }

            if (biggest !== index) {
                [this._data[biggest], this._data[index]] = [this._data[index], this._data[biggest]];
                index = biggest;
                continue;
            }

            break;
        }
    }

    size() {
        return this._size;
    }

}

思考二

利用快速排序的分治思想，无需完全排序数组，通过分区操作定位第k大元素，平均时间复杂度可达O(n)。

算法过程

1. 转换目标位置：第k个最大元素等价于数组升序排序后第n-k个位置的元素（n为数组长度）。
2. 分区操作：
   • 选择基准值（pivot），通过双指针将数组分为两部分：左侧元素≤基准，右侧元素≥基准。
   • 分区后，基准值的最终位置为j。
3. 递归查找：
   • 若目标位置k≤j，则在左半区间[l, j]继续查找。
   • 否则在右半区间[j+1, r]继续查找。
   • 当区间缩小到单个元素时，该元素即为结果。

时空复杂度

• 时间复杂度：平均O(n)，最坏O(n²)（可通过随机选择基准值优化至期望O(n)）。
• 空间复杂度：O(log n)，来自递归调用栈（最坏情况O(n)，对应极端不平衡分区）。

该算法通过"剪枝"避免了对无关区间的处理，相比全排序（O(n log n)）更高效，尤其适合大规模数据。

代码

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function(nums, k) {
    const n = nums.length;
    return quickSelect(nums, 0, n - 1, n - k);
};


function quickSelect(nums, l, r, k) {
    if (l === r) return nums[k];

    const pivot = nums[l];
    let i = l - 1, j = r + 1;
    while (i < j) {
        do {
            i++; 
        }
        while (nums[i] < pivot);
        do {
            j--; 
        } 
        while (nums[j] > pivot);

        if (i < j) {
            [nums[i], nums[j]] = [nums[j], nums[i]];
        }
    }
    if (k <= j) return quickSelect(nums, l, j, k);
    return quickSelect(nums, j + 1, r, k);   
}