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

数组中的第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$

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

时间复杂度$O（nlogk）$ 空间复杂度$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 手写堆维护

思路：

• $nums$种存放二叉堆，索引$[0,n-1]$对应按层序遍历对应的元素，对于下标从0开始的某节点$i$,左右孩子节点编号分别为$i\ast 2+1,i\ast 2+2$
• 下沉操作：$maxHeapify$操作为将二叉堆数组$nums$索引$i$处元素下沉
• 建堆操作：我们从最后一个非叶子节点$heapSize/2-1$开始倒序遍历，从下往上下沉
• 删除操作：每次将堆顶交换到数组末尾，再将$heapSize$减一，最后再调整新的堆顶即可
• 时间复杂度$O（nlogn）$ 空间复杂度$O（logn)$ 因为最大堆需要维护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);
}

• 快速选择（$Quickselect$）算法是一种用于在未排序的列表中找到第k小（或第k大）元素的高效算法。与快速排序一样，它使用分治策略，但不同于快速排序的是，它只递归处理包含目标元素的那一部分，而不是全部。这使得快速选择的平均时间复杂度为$O(n)$，最坏情况为$O(n^2)$，但在实际应用中，通过随机化选取枢轴（$pivot$）可以避免最坏情况。
• 复杂度：递归时，每层时间复杂度为$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;
}

参考

• LeetCode官方题解：数组中的第K个最大元素
• yxc常用代码模板1——基础算法
• LeetCode 215. 数组中的第K个最大元素