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

题目

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例：

  1. 输入: [3,2,1,5,6,4] 和 k = 2
  2. 输出: 5
  1. 输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
  2. 输出: 4

解答 & 代码

解法一：快速选择 O(N)

rand 头文件：#include <stdlib.h>

  1. class Solution {
  2. private:
  3.     int partition(vector<int>& nums, int left, int right)
  4.     {
  5.         int randomIdx = rand() % (right - left + 1) + left;
  6.         swap(nums[left], nums[randomIdx]);
  7.         int pivot = nums[left];
  8.         while(left < right)
  9.         {
  10.             while(left < right && nums[right] <= pivot)
  11.                 --right;
  12.             nums[left] = nums[right];
  13.             while(left < right && nums[left] >= pivot)
  14.                 ++left;
  15.             nums[right] = nums[left];
  16.         }
  17.         nums[left] = pivot;
  18.         return left;
  19.     }
  21.     void quickSelect(vector<int>& nums, int left, int right, int k)
  22.     {
  23.         if(left >= right)
  24.             return;
  26.         int pivot = partition(nums, left, right);
  27.         if(pivot == k - 1)
  28.             return;
  29.         else if(pivot > k - 1)
  30.             quickSelect(nums, left, pivot - 1, k);
  31.         else
  32.             quickSelect(nums, pivot + 1, right, k);
  33.     }
  34. public:
  35.     int findKthLargest(vector<int>& nums, int k) {
  36.         quickSelect(nums, 0, nums.size() - 1, k);
  37.         return nums[k - 1];
  38.     }
  39. };

复杂度分析：设数组长为 N

  • 时间复杂度 O(N)：粗略估计，假设每次扔掉一半，则
    • [中等] 215. 数组中的第K个最大元素 - 图1，其中 [中等] 215. 数组中的第K个最大元素 - 图2，即 [中等] 215. 数组中的第K个最大元素 - 图3，所以 [中等] 215. 数组中的第K个最大元素 - 图4
    • 等比数列求和公式：[中等] 215. 数组中的第K个最大元素 - 图5
    • 因此 [中等] 215. 数组中的第K个最大元素 - 图6
  • 空间复杂度 O(log N)：递归栈空间

执行结果：

  1. 执行结果：通过
  3. 执行用时：9 ms, 在所有 C++ 提交中击败了 72.53% 的用户
  4. 内存消耗：9.6 MB, 在所有 C++ 提交中击败了 89.72% 的用户

解法二：小根堆(优先队列) O(Nlogk)

  1. class Solution {
  2. public:
  3.     int findKthLargest(vector<int>& nums, int k) {
  4.         priority_queue<int, vector<int>, greater<int>> minHeap;
  5.         for(int i = 0; i < k; ++i)
  6.             minHeap.push(nums[i]);
  7.         for(int i = k; i < nums.size(); ++i)
  8.         {
  9.             if(nums[i] > minHeap.top())
  10.             {
  11.                 minHeap.pop();
  12.                 minHeap.push(nums[i]);
  13.             }
  14.         }
  15.         return minHeap.top();
  16.     }
  17. };

复杂度分析：设数组长为 N

  • 时间复杂度 O(N logk)：小根堆最多存储 k 个元素，因此每次查找、删除的时间复杂度 O(log k)，共 N 次
  • 空间复杂度 O(k)：小根堆最多存储 k 个元素

执行结果：

  1. 执行结果：通过
  3. 执行用时：16 ms, 在所有 C++ 提交中击败了 33.98% 的用户
  4. 内存消耗：9.9 MB, 在所有 C++ 提交中击败了 13.52% 的用户

解法三：小根堆（手造）

  1. class Solution {
  2. private:
  3.     void adjustHeap(vector<int>& nums, int start, int end)
  4.     {
  5.         int dadIdx = start;             // 父节点序号
  6.         int sonIdx = dadIdx * 2 + 1;    // 左子节点序号
  7.         while(sonIdx <= end)
  8.         {
  9.             if(sonIdx + 1 <= end && nums[sonIdx + 1] < nums[sonIdx])
  10.                 sonIdx += 1;
  11.             if(nums[dadIdx] < nums[sonIdx])
  12.                 return;
  13.             else
  14.             {
  15.                 swap(nums[dadIdx], nums[sonIdx]);
  16.                 dadIdx = sonIdx;
  17.                 sonIdx = dadIdx * 2 + 1;
  18.             }
  19.         }
  20.     }
  22.     void buildMinHeap(vector<int>& nums)
  23.     {
  24.         // 从最后一个非叶子节点开始调整
  25.         for(int i = nums.size() / 2 - 1; i >= 0; --i)
  26.             adjustHeap(nums, i, nums.size() - 1);
  27.     }
  28. public:
  29.     int findKthLargest(vector<int>& nums, int k) {
  30.         vector<int> minHeap;
  31.         for(int i = 0; i < k; ++i)
  32.             minHeap.push_back(nums[i]);
  33.         buildMinHeap(minHeap);
  35.         for(int i = k; i < nums.size(); ++i)
  36.         {
  37.             if(nums[i] > minHeap[0])
  38.             {
  39.                 minHeap[0] = nums[i];
  40.                 adjustHeap(minHeap, 0, k - 1);
  41.             }
  42.         }
  44.         return minHeap[0];
  45.     }
  46. };

复杂度分析：设数组长为 N

  • 时间复杂度 O(N logk)：小根堆最多存储 k 个元素，因此每次查找、删除的时间复杂度 O(log k)，共 N 次
  • 空间复杂度 O(k)：小根堆最多存储 k 个元素

执行结果：

执行结果：通过

执行用时：8 ms, 在所有 C++ 提交中击败了 72.53% 的用户
内存消耗：9.7 MB, 在所有 C++ 提交中击败了 79.68% 的用户

扩展：求 const top k：二分查找

如果传入的数组是 const vector<int> nums，也就是不能修改原数组的元素值，那么如何在 O(1) 的空间复杂度内实现查找第 k 大的元素？

class Solution {
private:
    /* 统计数组 nums 中大于等于 x 的元素个数，并返回 */
    int countGreaterThanX(vector<int> nums, int x)
    {
        int cnt = 0;
        for(int i = 0; i < nums.size(); ++i)
        {
            if(nums[i] >= x)
                ++cnt;
        }
        return cnt;
    }
public:
    int findKthLargest(vector<int>& nums, int k) {
        // 先统计数组中的最小值和最大值，
        int minNum = INT_MAX;
        int maxNum = INT_MIN;
        for(int i = 0; i < nums.size(); ++i)
        {
            minNum = min(minNum, nums[i]);
            maxNum = max(maxNum, nums[i]);
        }

        // 二分查找：minNum 就相当于二分查找的 left，maxNum 就相当于二分查找的 right
        while(minNum <= maxNum)
        {
            int midNum = minNum + (maxNum - minNum) / 2;
            // 统计数组中大于等于 midNum 的元素个数 cnt
            int cnt = countGreaterThanX(nums, midNum);
            // 若 cnt >= k，说明需要增大 midNum，使得大于等于 midNum 的元素个数减少
            // 因此，收缩左边界
            if(cnt >= k)
                minNum = midNum + 1;
            // 否则，说明需要减小 midNum，使得大于等于 midNum 的元素个数增大
            // 因此，收缩右边界
            else
                maxNum = midNum - 1;
        }
        return maxNum;
    }
};

复杂度分析：设数组长为 N

  • 时间复杂度 O(N logN)：二分查找的时间复杂度是 O(logN)，而对于每次二分的 midNum，调用 countGreaterThanX 查找数组 nums 中大于等于 midNum 的元素个数，时间复杂度为 O(N)
  • 空间复杂度 O(1)：

执行结果：

执行结果：通过

执行用时：8 ms, 在所有 C++ 提交中击败了 72.53% 的用户
内存消耗：11.8 MB, 在所有 C++ 提交中击败了 5.09% 的用户