leetcode215：数组中的第K个最大元素

题目链接

leetcode

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

C++ 代码

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <tuple>
using namespace std;

/*
    小根堆：堆的大小为 K，堆顶保存第 K 大的元素

    时间复杂度：O(NlogK)
        其中 N 是数组的长度，K 是要找的第 K 个最大元素的位置。
        遍历数组，时间复杂度是 O(N) 。
        每次插入操作和弹出操作，时间复杂度是 O(logK)。
        所以，总的时间复杂度是 O(NlogK)。
    空间复杂度：O(K)
        额外的空间，存储最大的 K 个元素。
*/
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> pq; // 小根堆

        for (const int& num : nums) {
            pq.push(num);
            if (pq.size() > k) {
                pq.pop();
            }
        }
        return pq.top();
    }
};

/*
    快速排序的选择方法（基于 Hoare 分区方案）

    每次划分数组后，判断中轴值 pivot 在划分后的位置。 
    若位置正好等于 k ，则直接返回 pivot;
    否则，判断 pivot 在左边还是右边，继续找第 k 个最大元素。

    时间复杂度：O(n)
        递归时，每层时间复杂度为 O(n)，但不是都进入左右两部分递归，
        仅进入一侧递归在平均情况下数组长度会减半。
        时间复杂度为 n + n/2 + n/4 + ... + 1 = O(n)。
    空间复杂度：O(logn)
        递归使用栈空间的大小是 O(logn)。
*/
class Solution_1 {
public:
    int findKthLargest(vector<int>& nums, int k) {
        return quickSelect(nums, 0, nums.size()-1, k-1);
    }

    int quickSelect(vector<int>& nums, int l, int r, int k) {
        if (l == r) {
            return nums[l];
        }

        // 随机选择基准值
        int pivotIndex = l + rand() % (r - l + 1);
        int pivot = nums[pivotIndex];

        // Hoare 分区，逆序排列
        int i = l, j = r;
        while (i <= j) {
            // 从左边找 <= pivot
            while (i <= j && nums[i] > pivot) {
                i++;
            }
            // 从右边找 >= pivot
            while (i <= j && nums[j] < pivot) {
                j--;
            }
            if (i <= j) {
                swap(nums[i++], nums[j--]);
            }
        }
        // nums[l...j] >= pivot
        // nums[i...r] <= pivot

        // 判断基准的位置
        if (k <= j) {
            return quickSelect(nums, l, j, k); // k 在左半部分
        } else if (k >= i) {
            return quickSelect(nums, i, r, k); // k 在右半部分
        } else {
            return nums[k]; // k 在中间，直接返回
        }
    }
};

// 辅助函数：打印数组
void printArray(const vector<int>& nums) {
    cout << "[";
    for (size_t i = 0; i < nums.size(); i++) {
        cout << nums[i];
        if (i != nums.size() - 1) cout << ",";
    }
    cout << "]";
}

int main() {
    Solution solution;
    vector<vector<int>> nums_cases = {
        {3,2,1,5,6,4},
        {3,2,3,1,2,4,5,5,6}
    };
    vector<int> k_cases = {
        2, 4
    };

    for (int i = 0; i < nums_cases.size(); i++) {
        auto& nums = nums_cases[i];
        int k = k_cases[i];

        cout << "Input: nums = ";
        printArray(nums);
        cout << ", k = " << k << endl;
        int result = solution.findKthLargest(nums, k);
        cout << "Output: " << result << endl;
    }

    return 0;
}