题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

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

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

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

示例 2:

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

思路

1）先遍历数组，把比当前值小的值，放在small数组里；大的值，放在big里；等于的放在equal里

2）如果big数组的长度大于k，说明k在big里，递归big数组进行查找k

3）如果len(nums)-len(small)<k：说明k在small数组里，递归small数组，但k的值需替换为：k-[len(nums)-len(small)]，需要去掉在big数组里的前几个，k需要更新下，更新到在small数组里的位置

4）如果正好等于当前基准值，则直接返回

class Solution(object):
    def quick_select(self, nums, k):
        # 随机选择基准数
        pivot = nums[-1]
        big, equal, small = [], [], []
        # 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
        for num in nums:
            if num > pivot:
                big.append(num)
            elif num < pivot:
                small.append(num)
            else:
                equal.append(num)
        if k <= len(big):
            # 第 k 大元素在 big 中，递归划分
            return self.quick_select(big, k)
        if len(nums) - len(small) < k:
            # 第 k 大元素在 small 中，递归划分
            return self.quick_select(small, k - len(nums) + len(small))
        # 第 k 大元素在 equal 中，直接返回 pivot
        return pivot

    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        return self.quick_select(nums, k)


if __name__ == '__main__':
    s = Solution()
    nums = [3,2,1,5,6,4]
    k = 2
    print(s.findKthLargest(nums, k))