4.2 快速選擇

215. Kth Largest Element in an Array

題目描述

在一個未排序的陣列中，找到第 $k$ 大的數字。

輸入輸出範例

輸入一個陣列和一個目標值 $k$，輸出第 $k$ 大的數字。題目保證一定有解。

Input: [3,2,1,5,6,4] and k = 2
Output: 5

題解

快速選擇 通常用於解決 k-th Element 問題，可以在平均 $O(n)$ 時間複雜度和 $O(1)$ 空間複雜度下完成求解。快速選擇的實現與快速排序相似，但只需要找到第 $k$ 大的中樞（pivot），不需要對中樞左右再進行排序。與快速排序一樣，快速選擇一般需要先將陣列打亂，否則最壞情況下的時間複雜度為 $O(n^2)$。

如果直接使用上述快速排序的代碼運行，可能會在 LeetCode 平台上接近超時。我們可以使用空間換取時間，直接儲存比中樞值小和大的元素，盡量避免進行交換操作。

C++

int findKthLargest(vector<int> nums, int k) {
    int pivot = rand() % nums.size();
    int pivot_val = nums[pivot];
    vector<int> larger, equal, smaller;
    for (int num : nums) {
        if (num > pivot_val) {
            larger.push_back(num);
        } else if (num < pivot_val) {
            smaller.push_back(num);
        } else {
            equal.push_back(num);
        }
    }
    if (k <= larger.size()) {
        return findKthLargest(larger, k);
    }
    if (k > larger.size() + equal.size()) {
        return findKthLargest(smaller, k - larger.size() - equal.size());
    }
    return pivot_val;
}

Python

def findKthLargest(nums: List[int], k: int) -> int:
    pivot_val = random.choice(nums)
    larger, equal, smaller = [], [], []
    for num in nums:
        if num > pivot_val:
            larger.append(num)
        elif num < pivot_val:
            smaller.append(num)
        else:
            equal.append(num)
    if k <= len(larger):
        return findKthLargest(larger, k)
    if k > len(larger) + len(equal):
        return findKthLargest(smaller, k - len(larger) - len(equal))
    return pivot_val

複雜度分析

• 時間複雜度:
  • 平均情況：$O(n)$，因為每次劃分會縮小問題規模。
  • 最壞情況：$O(n^2)$，當每次選擇的樞軸總是最差值（例如極大值或極小值）。
• 空間複雜度: $O(n)$，需要額外的空間來存放三個分組的元素。