LeetCode 215. Kth Largest Element in an Array

LeetCode 215. Kth Largest Element in an Array

這題也是 Top K Elements 的題型，直覺就是想到透過 Heap 來處理，這類型的題目八九不離十都可以使用 Heap 解決。

後續還會透過搭配其他資料結構來進行處理，我們來看一下原題內容講些什麼？

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

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

Example 2:

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

Constraints:

  • 1 <= k <= nums.length <= 10410^4
  • −104-10^4 <= nums[i] <= 10410^4

題目給你一組陣列，請你找出第 k 大的元素，這題型其實非常簡單，只要全部塞進 MaxHeap，因為 MaxHeap 每次塞進去都會重新排序一次，因此 Root 就是最大的元素，只要逐一將最大取出，取到第 k 個就代表最後取出來的元素就是答案。

照之前的做法先宣告一個 MaxHeap，將其所有元素塞進去，讓 MaxHeap 幫你排好。

val maxHeap = PriorityQueue<Int>{ n1, n2 ->
    n2 - n1
}
for(i in nums.indices){
    maxHeap.add(nums[i])
}

接著我們透過將其 Heap 內的元素透過 Heap.poll() 一個一個取出，直到第 k - 1 個，那我們就可以知道最後一個就是我們的答案。

以下是完整的解答。

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        val maxHeap = PriorityQueue<Int>{ n1, n2 ->
            n2 - n1
        }
        for(i in nums.indices){
            maxHeap.add(nums[i])
        }
        for(i in 1..k-1){
            maxHeap.poll()
        }
        return maxHeap.poll() ?: 0
    }
}

由於我們塞進 Heap 會做 N次，但是每次塞進去 Heap 會幫我們 LogN 的排序，這樣就會有 N 次的 LogN，所以整體複雜度就會是 O(NLogN)。

其實我們還可以改善這個作法，首先先將 k-1 個元素放進 MinHeap，想像一下最小的一定是在最上面。

for(i in 0..k-1){
    minHeap.add(nums[i])
}

接著我們把陣列第 k 個檢查元素是否比該 k 個元素大，可以透過 Heap.peek() 這個方法檢查，這個方法不會真的把元素從 Heap 取出，而是偷看一下它的元素值，如果真的比較大，我們才把元素取出，為什麼這樣做呢？

原因出在如果真的取出，整個 Heap 會再重新排序一次，所以如果只是單純想看頂部元素的值，就可以透過 Heap.peek() 這個方法來最佳化。

直到跑完所有陣列元素，我們 Heap 內裝的就會是前 k 大的元素集合，此時把最大的集合中的最小的元素取出，就是我們的答案了。

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

此時，我們將 k 個元素塞進 Heap，你會發現 heap 已經將其排序好。

minHeap:[2,3,5]  nums:[(3,2,5),4,1]

接著我們把剩餘的元素 4、1 都繼續塞進 Heap。

step 1 minHeap:[2,3,5]  nums:[(3,2,5),4,1]
step 2 minHeap:[3,4,5]  偷看 root 4 > 2 取出 2 加入 4
step 3 minHeap:[3,4,5]  偷看 root 1 < 3 不動作

此時我們的 Heap 就收集了原 nums 內前 k 大的元素，此時我們將 root 元素值也就是 3 取出即是答案。

以下是完整解答。

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        val minHeap = PriorityQueue<Int>{ n1, n2->
            n1 - n2
        }
        for(i in 0..k-1){
            minHeap.add(nums[i])
        }
        for(i in k..nums.size-1){
            if(nums[i] > minHeap.peek()){
                minHeap.poll()
                minHeap.add(nums[i])
            }
        }
        return minHeap.poll()?:0
    }
}

這樣的複雜度有稍微最佳化了一下，可以看出我們的複雜度一開始就是塞入了 k 個元素，接著就是進行了 n-k 次的排序，所以整體複雜度就是 O(KLogK+(N-K)LogK) 相當於 O(NLogK)。

其實還有更好的作法就是透過 QuickSelect，只要透過 QuickSelect 就可以在 O(N) 處理完畢，關於 QuickSelect 可以參考 https://magiclen.org/quickselect/ 。