寻找第K大

寻找第K大_牛客题霸_牛客网

问题简述

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

思路：快排

• 利用快排中的 partition 操作；

• partition 操作每次可以确定数组中某个数字在排序后的位置；

• 因为有一组测试用例会卡时间，所以需要用一些策略来决定 pivot，不能直接选第一个数字；

Python

class Solution:
    def findKth(self , a: List[int], n: int, K: int) -> int:
        # write code here
        import random
        
        def dfs(a, lo, hi):
            if lo >= hi: return
            
            # 法1）随机确定 pivot
            # r = random.randint(lo, hi)
            # a[lo], a[r] = a[r], a[lo]

            # 法2）三数取中
            mi = (lo + hi) // 2
            if a[lo] > a[mi]: a[lo], a[mi] = a[mi], a[lo]  # 确保 a[lo] < a[mi]
            if a[lo] > a[hi]: a[lo], a[hi] = a[hi], a[lo]  # 确保 a[lo] < a[hi]
            # 经过上面两步，a[mi]、a[hi] 中较小的就是中位数，启动到头部
            if a[mi] < a[hi]:
                a[lo], a[mi] = a[mi], a[lo]
            else:
                a[lo], a[hi] = a[hi], a[lo]
            
            p, l, r = lo, lo, hi
            while l < r:
                while l < r and a[r] <= a[p]:  # 因为是找第 K 大，所以要逆序排
                    r -= 1
                while l < r and a[l] >= a[p]:
                    l += 1
                a[l], a[r] = a[r], a[l]
            
            a[l], a[p] = a[p], a[l]
            
            # 第 K 大，排序后在数组中的索引为 K-1
            if l > K - 1: dfs(a, lo, l - 1)
            if l < K - 1: dfs(a, l + 1, hi)
            
        dfs(a, 0, n - 1)
        return a[K-1]