LeetCode面试题40. 最小的k个数(python,快速排序)

1. 题目

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

限制：

0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2. 代码

快排

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0 or len(arr) == 0:
            return []
        res = [0 for _ in range(k)]

        def select_base(res,begin,end):
            mid = (begin+end)//2
            if res[begin] > res[mid]:
                if res[begin] > res[end]:
                    if res[end] > res[mid]:
                        return end
                    else:
                        return mid
                else:
                    return begin
            else:
                if res[begin] < res[end]:
                    if res[mid] < res[end]:
                        return mid
                    else:
                        return end
                else:
                    return begin
        def swap(res, i, j):
            tmp = res[i]
            res[i] = res[j]
            res[j] = tmp

        def sort(res):
            quicksort(res,0,len(res)-1)

        def quicksort(res,begin,end):
            if begin>=end:
                return 
            left = begin
            right = end
            base = select_base(res,begin,end)
            swap(res,begin,base)
            while left < right:
                while left<right and res[right]>=res[begin]:
                    right -= 1
                while left<right and res[left]<=res[begin]:
                    left += 1
                if left<right:
                    swap(res,left,right)
                else:
                    break
            swap(res,begin,right)
            quicksort(res, begin, right-1 )
            quicksort(res, right+1, end)

        sort(arr)
                        
        return arr[:k]