【高频面试题】数组中的第K个最大元素（堆、快排进阶）

news2025/6/2 3:35:31

文章目录

- 数组中的第K个最大元素
- - 题目描述
  - 示例1
  - 示例2
  - 提示：
- 解法1（堆维护前k大元素）
- 解法2 手写堆维护
- 解法3（快速选择算法）
- 例题：P1923 【深基9.例4】求第 k 小的数
- 参考

数组中的第K个最大元素

题目描述

给定整数数组 $n u m s$ 和整数 $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$

解法1（堆维护前k大元素）

时间复杂度 $O （ n l o g k ）$ 空间复杂度 $O （ k)$

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> pq; 
        for(auto& num: nums){
            pq.emplace(num);
            if(pq.size() > k){
                pq.pop();   // 堆中元素超过k个，弹出最小的那个
            }
        }
        return pq.top();   
    }
};

解法2 手写堆维护

思路：

$n u m s$ 种存放二叉堆，索引 $[0, n - 1]$ 对应按层序遍历对应的元素，对于下标从0开始的某节点 $i$ ,左右孩子节点编号分别为 $i * 2 + 1, i * 2 + 2$
下沉操作： $ma x He a p i f y$ 操作为将二叉堆数组 $n u m s$ 索引 $i$ 处元素下沉
建堆操作：我们从最后一个非叶子节点 $h e a pS i ze /2 - 1$ 开始倒序遍历，从下往上下沉
删除操作：每次将堆顶交换到数组末尾，再将 $h e a pS i ze$ 减一，最后再调整新的堆顶即可
时间复杂度 $O （ n l o g n ）$ 空间复杂度 $O （ l o g n)$ 因为最大堆需要维护n个结点

class Solution {
public:
	// down 
    void maxHeapify(vector<int>& a, int i, int heapSize) {
        int l = i * 2 + 1, r = i * 2 + 2, largest = i;
        if (l < heapSize && a[l] > a[largest]) {
            largest = l;
        }
        if (r < heapSize && a[r] > a[largest]) {
            largest = r;
        }
        if (largest != i) {
            swap(a[i], a[largest]);
            maxHeapify(a, largest, heapSize);
        }
    }
    // initialize 
    void buildMaxHeap(vector<int>& a, int heapSize) {
        for (int i = heapSize / 2 - 1; i >= 0; i--) {
            maxHeapify(a, i, heapSize);
        }
    }
    int findKthLargest(vector<int>& nums, int k) {
        int heapSize = nums.size();
        buildMaxHeap(nums, heapSize); // 建堆
        // 执行k-1次提取最大值操作
        for (int i = nums.size() - 1; i >= nums.size() - k + 1; i--) {
            swap(nums[0], nums[i]);   // 调整堆顶
            -- heapSize;              // 删除堆
            maxHeapify(nums, 0, heapSize); 
        }
        return nums[0];
    }
};

解法3（快速选择算法）

我们首先来回顾一下快排的实现

void quickSort(vector<int>& nums, int l, int r) {
    if (l >= r) return;
    // - i 初始在左边界左侧（l-1）
    // - j 初始在右边界右侧（r+1）
    // - 基准值 x 选中间元素（避免极端情况如全排序数组导致最坏时间复杂度）
    int i = l - 1, j = r + 1;
    int x = nums[(l + r) >> 1]; // 位运算代替 (l + r) / 2，等价且更高效
    // partition ：双指针从两端向中间移动
    while (i < j) {
    	// 移动左指针 i：跳过所有小于 x 的元素，直到找到 >=x 的元素
        do i++; while (nums[i] < x); 
        // 移动右指针 j：跳过所有大于 x 的元素，直到找到 <=x 的元素
        do j--; while (nums[j] > x);
        //  如果指针未交错，交换左右指针的元素（确保左边 <=x，右边 >=x）
        if (i < j) swap(nums[i], nums[j]);
    }
    // 4. 递归排序左右子数组：
    quickSort(nums, l, j), quickSort(nums, j + 1, r);
}

快速选择（ $Q u i c k se l ec t$ ）算法是一种用于在未排序的列表中找到第k小（或第k大）元素的高效算法。与快速排序一样，它使用分治策略，但不同于快速排序的是，它只递归处理包含目标元素的那一部分，而不是全部。这使得快速选择的平均时间复杂度为 $O (n)$ ，最坏情况为 $O(n^2)$ ，但在实际应用中，通过随机化选取枢轴（ $p i v o t$ ）可以避免最坏情况。
复杂度：递归时，每层时间复杂度为 $O (n)$ ,但并不是都进入左右两部分递归。仅进入一侧递归在平均情况下 数组长度会减半，故平均情况下的时间复杂度为 $n + n /2 + n /4 + \dots + 1 = O (n)$

以下是快速选择实现找第K大元素的具体实现：

class Solution {
public:
    int quick_select(vector<int>& nums, int l, int r, int k) {
        if (l == r) return nums[k];
        int i = l - 1, j = r + 1, x = nums[l + r >> 1]; // 若选取nums[l], 极端样例 时间会很久
        //int x = nums[rand() % (r - l + 1) + l], i = l - 1, j = r + 1; // 随机选取
        while (i < j) {
            do i ++ ; while (nums[i] > x); // 注意是第k大 上面的模板是升序排序
            do j -- ; while (nums[j] < x);
            if (i < j) swap(nums[i], nums[j]);
        }
        if (k <= j) return quick_select(nums, l, j, k);
        else return quick_select(nums, j + 1, r, k);
    }

    int findKthLargest(vector<int>& nums, int k) {
        //srand(time(0)); // 随机种子
        return quick_select(nums, 0, nums.size() - 1, k - 1); // 0-base
    }
};

例题：P1923 【深基9.例4】求第 k 小的数

#include <bits/stdc++.h>

using namespace std;

const int N = 5e6 + 7;

int n, k;
int a[N];

int quick_select(int nums[], int l, int r, int k) {
    if (l == r) return nums[k];
    int x = nums[l], i = l - 1, j = r + 1;
    while (i < j) {
    	// paration 方向改一下即可
        do i ++; while (nums[i] < x); 
        do j --; while (nums[j] > x);
        if (i < j) swap(nums[i], nums[j]);
    }
    if (k <= j) return quick_select(nums, l, j, k);
    else return quick_select(nums, j + 1, r, k);
}

int main() {
    scanf("%d%d", &n, &k);
    //getchar();
    for (int i = 0; i < n; i++) {
        scanf("%d", &a[i]);
    }
    printf("%d\n", quick_select(a, 0, n - 1, k));
    return 0;
}