215. 第k大元素

Apr 25, 2024
1 views
Algorithm

题目说明

在未排序的数组中找到第 k 个最大的元素。请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

示例 1: 输入: [3,2,1,5,6,4] 和 k = 2 输出: 5

示例 2: 输入: [3,2,3,1,2,4,5,5,6] 和 k = 4 输出: 4

题解

使用快排的思想

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        k = n - k
        l, r = 0, n-1
        while l <= r:
            p = self.partition(nums, l, r)
            if p < k:
                l = p + 1
            elif p > k:
                r = p - 1
            else:
                return nums[p]

    def partition(self, nums, left, right):
        pivot = nums[left]
        l, r = left+1, right
        while l <= r:
            while l <= r and nums[l] < pivot:
                l += 1
            while l <= r and nums[r] >= pivot:
                r -= 1
            if l > r:
                break
            nums[l], nums[r] = nums[r], nums[l]
        nums[left], nums[r] = nums[r], nums[left]
        return r