剑指offer: 最小k个元素

思路：

大顶堆 – 最小k个元

• 求前k个最小用最大堆，求前k个最大用最小堆。
• 步骤如下：
  • 1: 取数组前k个元素初始化堆，从最后一个非叶子节点开始到根节点来构建大顶堆
  • 2: 当某个元素大于堆顶元素时，直接抛弃
  • 3: 当某个元素小于堆顶元素时，替换堆顶元素，再从堆顶重新构建大顶堆

代码（python3）

class Solution:
    @staticmethod
    def swap(heap, i, j):
        heap[i], heap[j] = heap[j], heap[i]

    def heapify(self, heap, i):
        left, right = 2 * i + 1, 2 * i + 2

        largest = i
        if left < len(heap) and heap[left] > heap[largest]:
            largest = left
        if right < len(heap) and heap[right] > heap[largest]:
            largest = right

        if largest != i:  # 调整当前子树为大顶堆
            self.swap(heap, i, largest)
            self.heapify(heap, largest)

    def build_heap(self, heap: list):
        # 从最后一个非叶子节点开始构建大根堆, len(heap)//2 - 1就是最后一个非叶子节点
        for i in range(len(heap) // 2 - 1, -1, -1):
            self.heapify(heap, i)

    def getLeastNumbers(self, arr: list, k: int):
        if not arr or k <= 0: return []
        if len(arr) <= k: return arr

        heap = arr[:k]
        self.build_heap(heap)  # 构建大根堆

        for i in range(k, len(arr)):
            if arr[i] < heap[0]:  # 当前元素比堆根小，则弹出堆根并插入当前元素
                heap[0] = arr[i]
                self.heapify(heap, 0)  # 重新调整堆根
        return heap