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)
