【LeetCode刷题】数组中的第K个最大元素

题目

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

code

• 随机选择一个枢轴。
• 使用划分算法将枢轴放在数组中的合适位置 pos。将小于枢轴的元素移到左边，大于等于枢轴的元素移到右边。
• 比较 pos 和 N - k 以决定在哪边继续递归处理。

在这里，由于知道要找的第 N - k 小的元素在哪部分中，我们不需要对两部分都做处理，这样就将平均时间复杂度下降到 O(N)。

import random

class Solution:
    def findKthLargest(self, nums, k):
        def partition(left, right, pivot_index):
            pivot = nums[pivot_index]
            # move pivot to end
            nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

            # move all smaller elements to the left
            store_index = left
            for i in range(left, right):
                if nums[i] < pivot:
                    nums[store_index], nums[i] = nums[i], nums[store_index]
                    store_index += 1

            # move pivot to its final place
            nums[right], nums[store_index] = nums[store_index], nums[right]

            return store_index

        def select(left, right, k_smallest):

            # if the list contains only one element,return that element
            if left == right:
                return nums[left]

            # select a random pivot_index between
            pivot_index = random.randint(left, right)

            # find the pivot position in a sorted list
            pivot_index = partition(left, right, pivot_index)

            # the pivot is in its final sorted position
            if k_smallest == pivot_index:
                return nums[k_smallest]

            # go left
            elif k_smallest < pivot_index:
                return select(left, pivot_index - 1, k_smallest)
            # go right
            else:
                return select(pivot_index + 1, right, k_smallest)

        # kth largest is (n-k)th smallest
        return select(0, len(nums) - 1, len(nums) - k)


a = Solution()
kth_largest = a.findKthLargest(nums=[3,2,1,5,6,4],k=2)
print(kth_largest)