牛客网——寻找第K大

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

给定一个整数数组a,同时给定它的大小n和要找的K(K在1到n之间),请返回第K大的数,保证答案存在。

快速排序:快速排序算法首先会从列表里选择一个被称为基准的元素,然后算法会把所有小于基准的元素移动到列表的左侧,把所有大于基准的元素移动到列表的右侧。之后,基准元素会被放在列表的两部分之间,这个位置就会是整个列表排序之后基准元素应该在的正确位置。下一步就是递归地对基准元素两侧的两个较小的数组再次进行上面那样的排序。 快速排序的代码:

import numpy
def quicksort(a,left,right):
    if left < right:
        pivot = a[right]         #基准元素选择最后一个元素
        i = left
        j = right - 1
        while True:
            while a[i] > pivot and i <= right:   #比基准元素大的在左边
                i += 1
            while a[j] < pivot and j >= left:    #比基准元素小的在右边
                j -= 1
            if i < j :
                a[i],a[j] = a[j],a[i]   #a[i]是从左至右第一个比基准元素小的元素,a[j]是从右至左第一个比基准元素大的元素
            else:
                break
        a[i],a[right] = a[right],a[i]  #a[i]是从左至右第一个比基准元素小的元素,并且a[i]的右侧也是比基准元素小的元素(除最后一个基准元素除外)所以基准元素应该在i位置
        quicksort(a,left,i-1)
        quicksort(a,i+1,right)
a = [8,9,4,1,5,6,2,3]
quicksort(a,0,len(a)-1)
print(a)

上面快速排序的代码实现从大到小排列,如果需要从小到大排序只需改变代码中的大于小于号。 在快速排序的基础上找到第K大,通过判断基准元素的位置,只需进行一半的操作,如果基准元素的位置就是K,则返回基准元素,如果基准元素的位置大于K,则进行左半区域搜索,否则,进行右半区域搜索。

class Solution:
    def findKth(self,a,n,K):
        res = []
        def dfs(left,right):
            if left < right:
                pivot = a[right]
                i = left
                j = right - 1
                while True:
                    while a[i] > pivot and i <= right:
                        i += 1
                    while a[j] < pivot and j >= left:
                        j -= 1
                    if i < j:
                        a[i], a[j] = a[j], a[i]
                    else:
                        break
                a[i], a[right] = a[right], a[i]
                if K - 1 == i:
                    res.append(a[i])
                    return None
                elif K-1 < i:
                    dfs(left, i - 1)
                else:
                    dfs(i + 1, right)
            else:
                res.append(a[right])
                return None
        dfs(0, len(a) - 1)
        return res[0]
    a = [8,9,4,1,5,6,2,3]
    n = 6
    K = 2
    print(findKth("", a, n, K))
经验分享 程序员 微信小程序 职场和发展