python--剑指offer--40. 最小的k个数

import heapq
from typing import List


class Solution:

    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return list()

        hp = [-x for x in arr[:k]]
        heapq.heapify(hp)
        for i in range(k, len(arr)):
            print(hp)
            if -hp[0] > arr[i]:
                heapq.heappop(hp)
                heapq.heappush(hp, -arr[i])
        ans = [-x for x in hp]
        return ans


if __name__ == '__main__':
    solution = Solution()
    arr = [3, 2, 1]
    k = 2
    res = solution.getLeastNumbers(arr, k)
    print(res)

from typing import List


class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
        self.randomized_selected(arr, 0, len(arr) - 1, k)
        return arr[:k]

    def randomized_selected(self, arr, l, r, k):
        pos = self.randomized_partation(arr, l, r)
        num = pos - l + 1
        if num < k:
            self.randomized_selected(arr, pos+1, r, k-num)
        elif num > k:
            self.randomized_selected(arr, l, pos-1, k)

    def randomized_partation(self, arr, l, r):
        import random

        i = random.randint(l, r)
        arr[i], arr[r] = arr[r], arr[i]
        return self.partation(arr, l, r)

    def partation(self, arr, l, r):
        pos = l
        for i in range(l, r):
            if arr[i] < arr[r]:
                arr[i], arr[pos] = arr[pos], arr[i]
                pos += 1

        arr[pos], arr[r] = arr[r], arr[pos]
        return pos


if __name__ == '__main__':
    solution = Solution()
    arr = [3, 2, 1]
    k = 2
    res = solution.getLeastNumbers(arr, k)
    print(res)