leetcode 寻找第 K 大(快排)

描述
有一个整数数组，请你根据快速排序的思路，找出数组中第 k 大的数。

给定一个整数数组 a ,同时给定它的大小n和要找的 k ，请返回第 k 大的数(包括重复的元素，不用去重)，保证答案存在。
要求：时间复杂度 O(nlogn)，空间复杂度 O(1)
数据范围：0≤n≤pow(10,5) ， 1≤K≤n，数组中每个元素满足 0≤val≤10000000

示例1
输入：[1,3,5,2,2],5,3
返回值：2

基本思路还是快排的思路， 注意不要在
if i < j:
self.swap(a, i, j)
if 语句里面加一些 i +=1, j-=1 之类的，逻辑就不对了。
如 13，11，5，6，12，13，5 这种情况， 另外，对于极端已排好序的情况，最差时间复杂度 是 O(n*n)，有一些 case 可能过不了，可以增加一个随机的 shuffle

# -*- coding:utf-8 -*-
import random

class Solution:
    def swap(self, a, i, j):
        temp = a[i]
        a[i] = a[j]
        a[j] = temp
        
    def findK(self, a, l, r, K):
        if l == r:
            return a[l]
        pivot = a[l]
        i, j = l, r
        while i<j:
            while i<j and a[j]>=pivot:
                j -= 1
            while j>i and a[i]<=pivot:
                i += 1          
            if i < j:
                self.swap(a, i, j)
        if a[j] < a[l]:
            self.swap(a, j, l)
        if j-l+1 == K:
            return a[j]
        elif j-l+1 > K:
            return self.findK(a, l, j-1, K)
        else:
            return self.findK(a, j+1, r, K-j+l-1)
        
        
    def findKth(self, a, n, K):
        # write code here
        random.shuffle(a)
        return self.findK(a, 0, n-1, n+1-K)