用堆找出最小的 N 个数

import random


class HeapForMinN:
    """求最小的 N 个数"""

    def __init__(self, size):
        self.size = size
        self.heap = []
    
    def add(self, num):
        if len(self.heap) < self.size:
            # 数据太少
            self.heap.append(num)
        elif len(self.heap) + 1 == self.size:
            # 数据数量够了，构建堆
            self.heap.append(num)
            self.build_heap()
        elif num < self.heap[0]:
            # 使用新数替换堆顶
            self.heap[0] = num
            self.max_heapify(0)

    def build_heap(self):
        length = len(self.heap)
        for i in range((length + 1) // 2, -1, -1):
            self.max_heapify(i)

    def max_heapify(self, index):
        length = len(self.heap)
        while index < length:
            maxi = index
            i1, i2 = index * 2 + 1, index * 2 + 2
            if i1 < length and self.heap[maxi] < self.heap[i1]:
                maxi = i1
            if i2 < length and self.heap[maxi] < self.heap[i2]:
                maxi = i2
            if maxi == index:
                break
            self.heap[index], self.heap[maxi] = self.heap[maxi], self.heap[index]
            index = maxi

    def min_n(self):
        return self.heap


if __name__ == "__main__":
    list_ = list(range(1, 10001)) * 2
    random.shuffle(list_)
    N = 10
    hn = HeapForMinN(N)
    for i in list_:
        hn.add(i)
    min_n = hn.min_n()
    print(min_n)        # [5, 4, 5, 4, 3, 3, 1, 1, 2, 2]
    assert sorted(min_n) == sorted(list_)[:n]