面试题40. 最小的k个数

面试题40. 最小的k个数

这里可以使用随机选择算法，平均时间复杂度为 O（n），最差情况下时间复杂度为 O（n^2)。另外还可以使用更加复杂 BFPRT 算法 好处是平均时间复杂度和最差情况下时间复杂度都是 O（n）。

import random
class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if not arr or k == 0:
            return []
        def partition(l, r):
            pivot = arr[r]
            j = l-1
            for i in range(l, r):
                if arr[i] <= pivot:
                    j += 1
                    arr[i], arr[j] = arr[j], arr[i]
            arr[j+1], arr[r] = arr[r], arr[j+1]
            return j+1
        
        def topk(l, r, k):
            i = random.randint(l, r)
            arr[r], arr[i] = arr[r], arr[i]
            m = partition(l, r)
            if m+1 == k:
                return
            if m+1 < k:
                topk(m+1, r, k)
            else:
                topk(l, m-1, k)
                
        topk(0, len(arr)-1, k)
        return arr[:k]