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

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

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:
nput: nums = [3,2,1,5,6,4], k = 2
Output: 5

示例 2:
Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

解题思路

输入： 整数数组 nums 和整数 k

输出： 数组中第 k 大的元素

方法分析

• 直接排序的时间复杂度是 O(n log n)，不符合题目希望的更优解法。
• 我们可以使用 最小堆（Min Heap） 来降低复杂度。

具体做法

1. 实现一个最小堆 MiniHeap。

2. 遍历数组元素：

   • 如果堆的大小小于 k，直接入堆。
   • 如果堆的大小等于 k 且当前元素大于堆顶（最小值），则用该元素替换堆顶并下沉。

3. 遍历结束后，堆中保留的就是 最大的 k 个元素。

4. 堆顶元素即为第 k 大的数。

复杂度分析

• 每次堆操作为 O(log k)，共处理 n 个元素。
• 总时间复杂度：O(n log k)
• 空间复杂度：O(k)

代码实现

python

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 将输入的列表 nums 转换成最小堆
        # heapq.heapify 函数会原地对列表进行堆化处理
        heapq.heapify(nums)

        # 不断弹出堆顶元素，直到堆中剩下 k 个元素
        # 因为堆是最小堆，所以最终剩下的堆顶元素即为第 k 大的元素
        while len(nums) > k:
            heapq.heappop(nums)
        
        # 返回堆顶的元素，即为第 k 大的元素
        return nums[0]

javascript

// 定义一个最小堆类 MiniHeap
class MiniHeap {
    constructor() {
        this.data = [];  // 用数组存储堆，完全二叉树的顺序结构
    }

    // 返回堆顶元素（最小值）
    peek() { return this.data[0]; }

    // 返回堆中元素个数
    size() { return this.data.length; }

    // 插入一个元素
    push(val) {
        this.data.push(val);                 // 新元素放在数组末尾（树的最右下角）
        this._siftUp(this.size() - 1);       // 从下往上调整，恢复最小堆性质
    }

    // 删除并返回堆顶（最小值）
    pop() {
        if (this.size() === 0) return null;  // 空堆直接返回 null

        const top = this.data[0];            // 记录堆顶
        const last = this.data.pop();        // 删除数组最后一个元素

        if (this.size() > 0) {
            this.data[0] = last;             // 把最后一个元素放到堆顶
            this._siftDown(0);               // 从上往下调整，恢复堆序
        }

        return top;                          // 返回最小值
    }

    // 替换堆顶（效率比先 pop 再 push 更高）
    replaceTop(val) {
        this.data[0] = val;                  // 直接覆盖堆顶
        this._siftDown(0);                   // 下沉恢复堆序
    }

    // 上浮操作：插入时用
    _siftUp(i) {
        while (i > 0) {
            const parent = Math.floor((i - 1) / 2); // 父节点下标
            if (this.data[i] >= this.data[parent]) break; // 已满足最小堆，不需要调整

            // 否则和父节点交换，继续往上走
            [this.data[i], this.data[parent]] = [this.data[parent], this.data[i]];
            i = parent;
        }
    }

    // 下沉操作：删除堆顶/替换堆顶时用
    _siftDown(i) {
        const n = this.size();
        while (true) {
            let smallest = i;               // 先假设当前节点是最小的
            const left = i * 2 + 1;         // 左子节点下标
            const right = left + 1;         // 右子节点下标

            // 如果左子节点更小，更新 smallest
            if (left < n && this.data[left] < this.data[smallest]) smallest = left;
            // 如果右子节点更小，更新 smallest
            if (right < n && this.data[right] < this.data[smallest]) smallest = right;

            if (smallest === i) break;      // 已经满足最小堆，不用调整

            // 否则交换，并继续向下
            [this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
            i = smallest;
        }
    }
}


// 主函数：找数组中的第 k 大元素
var findKthLargest = function(nums, k) {
    const heap = new MiniHeap();  // 定义一个最小堆

    // 遍历数组
    for (let x of nums) {
        if (heap.size() < k) {
            // 堆的大小不足 k，就直接放进去
            heap.push(x);
        } else {
            // 如果堆已满，并且当前元素比堆顶大
            // 说明它更“有资格”进入前 k 大
            if (x > heap.peek()) {
                heap.replaceTop(x);  // 替换掉堆顶最小值
            }
        }
    }

    // 遍历完成后，堆中只保留了最大的 k 个数
    // 堆顶就是其中最小的那个，也就是“第 k 大”
    return heap.peek();
};

复杂度分析

时间复杂度：O(n)

空间复杂度：O(1)

链接

215 国际版

215 中文版