快速选择算法

2019-01-05

数组中的第K个最大元素

数组中的第K个最大元素

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

  • 可以用先排序再取值的方法来做，这样做最好的时间复杂度为O(NlogN)。
  • 可以使用快速排序的一个变种quick_select，则平均复杂度为O(N)。

通过一趟快排过后，序列将被分为(比key小的数,key,比key大的数)三部分。那么假设key的下标为i，如果k < i，则第K大的数必然在快排左边的区域；如果k = i，则key就是第k大的数；如果k > i，则k必然在快排的右边的区域。接下来递归即可得到第k大的数。

def helper(nums, k, left, right):
    if left == right:
        # 只有1个数，那么就是这个数
        return nums[left]
    if left < right:
        i = left
        j = right
        key = nums[left]
        while i < j:
            # j向前搜索，找到第一个小于key的值
            while i < j and nums[j] < key:
                j -= 1
            if i < j:
                nums[i] = nums[j]
                i += 1
            # i向后搜索，找到第一个大于key的值
            while i < j and nums[i] >= key:
                i += 1
            if i < j:
                nums[j] = nums[i]
                j -= 1
        nums[i] = key
        # 递归左边部分和右边部分
        if left < i and k < i:
            return helper(nums, k, left, i - 1)
        if right > i and k > i:
            return helper(nums, k, i + 1, right)
        return nums[i]

# 获得第K大的数
def getTopk(nums, k):
    # 第k大的数下标应该是k-1
    return helper(nums, k-1, 0, len(nums)-1)

if __name__ == '__main__':
    nums = [3,2,1,5,6,4]
    k = 2
    print(getTopk(nums, k))
    
输出：
5