ssutsen / DS.

111-2師大科技系資料結構
0 stars 0 forks source link

215. Kth Largest Element in an Array #13

Open ssutsen opened 1 year ago

ssutsen commented 1 year ago

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.

You must solve it in O(n) time complexity.

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

Link

ssutsen commented 1 year ago

quick selection solution

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        k = len(nums) - k

        def quickSelect(l, r):
            pivot, p = nums[r], l
            for i in range(l,r):
                if nums[i] <= pivot:
                    nums[p],nums[i] = nums[i], nums[p]
                    p += 1
            nums[p], nums[r] = nums[r] , nums[p]

            if p > k: return quickSelect(l, p-1)
            elif p < k: return quickSelect(p+1, r)
            else: return nums[p]
        return quickSelect(0, len(nums) - 1)
ssutsen commented 1 year ago

heap solution

class Solution(object):
    def findKthLargest(self, nums, k):
        min_heap = []
        for num in nums:
            heapq.heappush(min_heap, num)
            if len(min_heap) > k:
                heapq.heappop(min_heap)
        return min_heap[0]

首先,建立一個空的小根堆 min_heap,然後依次遍歷 nums 中的每個元素 num。

對於每個 num,將其加入小根堆中,然後檢查小根堆的大小是否超過 k。如果超過了 k,就從小根堆中彈出一個元素,因為我們只需要維護最大的 k 個元素。注意,因為我們是要求第 k 大元素,因此我們需要維護一個大小為 k 的小根堆,堆頂元素就是第 k 大元素。

最後,返回小根堆的堆頂元素,即第 k 大元素。