数组中的第K个最大元素

题意

找出数组中第 k 大的元素（排序后倒数第 k 个）。

思路

维护一个大小为 k 的最小堆。每次插入元素后，若堆大小超过 k，弹出堆顶最小值。最终堆中保留最大的 k 个元素，堆顶就是第 k 大。

最小堆而非最大堆的关键理解：我们要保留最大的 k 个元素，所以每次弹出的是堆中最小的那个——即这 k 个中"最不够格"的。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    int findKthLargest(vector<int> &nums, int k) {
        priority_queue<int, vector<int>, greater<>> pq;
        for (int x : nums) {
            pq.push(x);
            if ((int)pq.size() > k)
                pq.pop();
        }
        return pq.top();
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, k;
    cin >> n >> k;
    vector<int> a(n);
    for (int &x : a)
        cin >> x;
    cout << Solution().findKthLargest(a, k) << '\n';
    return 0;
}

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23#!/usr/bin/env python3
import heapq
from typing import List


class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        pq = []
        for x in nums:
            heapq.heappush(pq, x)
            if len(pq) > k:
                heapq.heappop(pq)
        return pq[0]


def main():
    n, k = map(int, input().split())
    a = list(map(int, input().split()))
    print(Solution().findKthLargest(a, k))


if __name__ == "__main__":
    main()

复杂度

时间复杂度： $O(n \log k)$ ，每个元素堆操作 $O(\log k)$ 。
空间复杂度： $O(k)$ ，堆最多存 k 个元素。

总结

第 k 大/小元素用大小为 k 的堆：第 k 大用最小堆（弹小留大），第 k 小用最大堆（弹大留小）。堆顶即答案。

数组中的第K个最大元素

目录

题意

思路

代码

复杂度

总结