215. Kth Largest Element in an Array

題目網址：https://leetcode.cn/problems/kth-largest-element-in-an-array/

題意：給一 array nums 和一整數 k, 返回第 k 大的元素。

注意：你需要找的是 nums 排序後的第 k 大的元素, 而非第 k 個不同的元素。

設計滿足 $O(n)$ time 的演算法。

Solution 1：

想法：利用 Heap, 維護前 k 大的元素, 類似 347. Top K Frequent Elements

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<>> pq; // min heap

        for (const auto& num : nums) {
            pq.emplace(num);

            // 若加入後 heap 中有 k + 1 個, 則 pop 掉最小的來維持 k 個
            if (pq.size() > k) {
                pq.pop();
            }
        }
        return pq.top();
    }
};

• time：$O(n \cdot log(k))$ ➔ for loop 最多 insert/delete n 個點, 每次操作需 $O(log(k))$, 因為 pq 中最多 k 個點
• space：$O(k)$ ➔ pq 的最大長度為 k

Solution 2：

想法：利用 Quick Select, 原題目就會變成求將 nums 由小到大排序後 idx = n - k 的元素, 因為排序完後 idx = n - 1 為第一大的元素, 故 idx = n - k 為第 k 大的元素

• 先選一元素當作 pivot
  • 把 val ≤ pivot 的元素放在 pivot 的左邊
  • 而 val > pivot 的元素放在 pivot 的右邊
• 比較 pivot_idx 是否為 k
  • 若 k > pivot_idx：遞迴搜索右邊的部分 quickSelect(pivot_idx + 1, right)
  • 若 k < pivot_idx：遞迴搜索左邊的部分 quickSelect(left, pivot_idx - 1)
  • 若 k == pivot_idx：返回 nums[k]

e.g. nums = [3,2,1,5,6,4], k = 2 ➔ k = n - k = 6 - 2 = 4

• 首先, left = 0, right = 5, 取出 pivotVal = 4, pivotIdx = 0,
• 第一輪 Quick Select 做完, 得到 nums = [3,2,1,4,6,5], pivotIdx = 3
• 由於 k = 4 > pivotIdx = 3, 故往右邊尋找 quickSelect(nums, pivotIdx + 1, right, k)
• 第二輪 Quick Select 做完, 得到 nums = [3,2,1,4,5,6], pivotIdx = 4
• k == pivotIdx, 故返回 nums[4] = 5

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        const int n = nums.size();
        k = n - k; // 若由小排到大, 則 nums 中 idx = n - k 的元素為第 k 大的元素
        return quickSelect(nums, 0, n - 1, k);
    }

private:
    int quickSelect(vector<int>& nums, int left, int right, int k){
        int pivot = nums[right], pivotIdx = left; // 取最右的元素為 pivot

        // 將比 pivot 小的元素從 idx = left 的位置開始放
        for (int i = left; i < right; ++i) {
            if (nums[i] <= pivot) {
                swap(nums[i], nums[pivotIdx++]);
            }
        }

        // 將 pivot 放在 pivotIdx 上, 使得左邊皆 <= pivot, 右邊皆 > pivot
        // 必須寫 nums[right], 不能寫 pivot, 否則 nums[pivotIdx] 仍會是 pivot
        swap(nums[pivotIdx], nums[right]);

        if (k < pivotIdx) { // 搜索左邊的部分
            return quickSelect(nums, left, pivotIdx - 1, k);
        } else if (k > pivotIdx) { // 搜索右邊的部分
            return quickSelect(nums, pivotIdx + 1, right, k);
        }
        
        return nums[pivotIdx]; // k == pivotIdx
    }
};

• time：$O(n)$ ➔ Quick Select 的 Avg case
  • worse case 為 $O(n^2)$, 也就是每次取完 pivot 都只篩掉一個元素
• space：$O(1)$ ➔ 只需常數空間