Leetcode 215. Kth Largest Element in an Array

发表于 | 分类于

问题描述

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

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

Example 2:

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

Note:

  • You may assume k is always valid, 1 ≤ k ≤ array's length.

解题思路

堆排序，建堆O(logn)，取前k个，总时间复杂度O(klogn)。

1
2
3
4
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        import heapq
        return heapq.nlargest(k, nums)[-1]

进阶

快排方法，将比pivot小的移动到pivot左边大的移动到右边，如果结束后pivot位于第k个则pivot即为答案。如果大于k则从pivot左边找，小于则从右边找。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return self.partition( 0, len(nums) - 1, nums, len(nums) - k)

    def partition(self, l, r, nums, k):

        pivoti = random.randint(l, r)
        pivot = nums[pivoti]
        nums[pivoti], nums[r] = nums[r], nums[pivoti]

        L = l
        R = r - 1

        while L <= R:
            while nums[L] < pivot:
                L += 1
            while nums[R] > pivot:
                R -= 1

            if L <= R:
                nums[L], nums[R] = nums[R], nums[L]
                L += 1
                R -= 1

        nums[L], nums[r] = nums[r], nums[L]

        if L == k:
            return nums[L]
        elif L > k:
            return self.partition(l, L - 1, nums, k)
        elif L < k:
            return self.partition(L + 1, r,  nums, k)