【LeetCode-數組】數組中的第K個最大元素

題目描述

在未排序的數組中找到第 k 個最大的元素。請注意，你需要找的是數組排序后的第 k 個最大的元素，而不是第 k 個不同的元素。
示例：

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

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

說明：

• 你可以假設 k 總是有效的，且 1 ≤ k ≤ 數組的長度。

題目鏈接： https://leetcode-cn.com/problems/kth-largest-element-in-an-array/

思路1

升序排序，然后輸出第 n-k 個元素（第 k 大）。代碼如下：

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        return nums[nums.size()-k];
    }
};

• 時間復雜度：O(nlogn)
  快速排序的時間復雜度。
• 空間復雜度：O(1)

思路2

使用小根堆，小根堆的堆頂元素是堆中最小的元素。所以，我們遍歷數組，將數組中的元素加入堆中，然后保持堆的大小為 k，這樣數組遍歷結束后的堆中的元素就是數組中前 k 大的元素，堆頂元素是堆中最小元素，也就是第 k 大的元素。代碼如下：

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> q;
        for(int i=0; i<nums.size(); i++){
            q.push(nums[i]);
            if(q.size()>k) q.pop();
        }
        return q.top();
    }
};

• 時間復雜度：O(nlogk)
  向堆中添加一個元素的時間復雜度為 logk，添加 n 次，所以為 nlogk。
• 空間復雜度：O(k)

思路3

使用快速排序中的 partition 過程來做。partition 過程如下：

• 找一個比較的基准，例如選取數組的第一個元素；
• 設置兩個指針 i=0，j=nums.size()-1；
• 如果 i<j，循環：
  • 從數組尾從后向前遍歷，找到第一個小於基准的元素 nums[i]；
  • 從數組頭從前向后遍歷，找到第一個大於基准的元素 nums[j]；
  • 交換 nums[i] 和 nums[j]；
• 交換 nums[i] 和 nums[left]。

一次 partition 之后，nums[i] 之前的數字都小於等於 nums[i]，nums[i] 之后的數字都大於等於 nums[i]，nums[i] 的位置已經確定，也就是 nums[i] 是數組 nums 中第 i 小的元素，如果 i==nums.size()-k（第 k 大元素），則返回 nums[i]，否則根據 i 和 nums.size()-k 的大小關系，使用類似於二分查找的方法縮小 partition 的范圍。具體代碼如下：

class Solution {
public:
    int ans = -1;
    int findKthLargest(vector<int>& nums, int k) {
        int left = 0;
        int right = nums.size()-1;
        int target = nums.size()-k;
        while(true){    // while(left<=right) 也行
            int idx = partition(nums, left, right);
            if(idx==target) return nums[idx];
            else if(idx<target){
                left = idx+1;
            }else if(idx>target){
                right = idx-1;
            }
        }
        return -1;
    }

    int partition(vector<int>& nums, int left, int right){
        int base = nums[left];
        int i = left;
        int j = right;
        while(i<j){
            while(nums[j]>=base && i<j) j--;
            while(nums[i]<=base && i<j) i++;
            swap(nums[i], nums[j]);
        }
        swap(nums[i], nums[left]);
        return i;
    }
};

這種寫法參考了這篇題解。

• 時間復雜度：O(n)
• 空間復雜度：O(1)