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]