215. Kth Largest Element in an Array

Medium

Description

給定一個整數陣列 nums 和整數 k，回傳陣列中第 k 大的元素。注意是排序後第 k 大，不是第 k 個不同的元素。要求設計不排序的演算法。

Example:

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

Intuition#

核心思路：第 K 大問題的兩大武器——Min Heap (O(n log k)) 或 Quick Select (平均 O(n))。

第 k 大 = 排序後 index 為 n - k 的元素
Min Heap 大小 k：遍歷完後堆頂即答案
Quick Select：類似 Quick Sort 的 partition，但只遞迴一邊，平均 O(n)

Approaches#

1: Min Heap#

概念: 維護大小為 k 的 Min Heap，遍歷所有元素，堆頂即第 k 大
時間複雜度: O(n log k)
空間複雜度: O(k)

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        val minHeap = PriorityQueue<Int>()

        for (num in nums) {
            minHeap.offer(num)
            if (minHeap.size > k) {
                minHeap.poll()
            }
        }

        return minHeap.peek()
    }
}

⭐2: Quick Select#

概念: 隨機選 pivot 做 partition，根據 pivot 位置決定只往左或右遞迴，平均只需 O(n)
時間複雜度: 平均 O(n)，最壞 O(n^2)（加隨機化可降低最壞情況機率）
空間複雜度: O(1)

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        val targetIdx = nums.size - k
        return quickSelect(nums, 0, nums.size - 1, targetIdx)
    }

    private fun quickSelect(nums: IntArray, left: Int, right: Int, target: Int): Int {
        val pivotIdx = left + (Math.random() * (right - left + 1)).toInt()
        val pivot = nums[pivotIdx]
        nums[pivotIdx] = nums[right].also { nums[right] = nums[pivotIdx] }

        var i = left
        for (j in left until right) {
            if (nums[j] <= pivot) {
                nums[i] = nums[j].also { nums[j] = nums[i] }
                i++
            }
        }
        nums[i] = nums[right].also { nums[right] = nums[i] }

        return when {
            i == target -> nums[i]
            i < target -> quickSelect(nums, i + 1, right, target)
            else -> quickSelect(nums, left, i - 1, target)
        }
    }
}

補充：排序法

最直觀的做法，排序後直接取。

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        nums.sort()
        return nums[nums.size - k]
    }
}

時間複雜度: O(n log n)
空間複雜度: O(1) 或 O(n) 取決於排序實作

🔑 Takeaways#

Pattern: 第 K 大/小問題，Heap 是穩定解，Quick Select 是最佳平均解
關鍵技巧: Quick Select 只遞迴一邊是效率關鍵；隨機選 pivot 可有效避免最壞情況