题目描述
给定整数数组 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))