heaps_running_median.py 文件源码

def add_number_to_heap(num):
    """Adds a number to a heap, assuming the heaps correctly represent the lower half of numbers and the upper half.
    So the only factor taken into account is the number of items in the heaps.
    If same number of items add to lower half, else add to half with the least number of items.
    """
    if len(lower_half_max_heap) == 0:
        heapq.heappush(lower_half_max_heap, num * -1)
    elif len(upper_half_min_heap) == 0:
        heapq.heappush(upper_half_min_heap, num)
    else:
        # Assume the heap is a valid split of the data.
        # Difference will be positive if upper_half_min_heap has more elements, else negative.
        difference = len(upper_half_min_heap) - len(lower_half_max_heap)
        if difference >= 1:
            # upper_half_min_heap has more elements, should put this next number in lower_half_max_heap.
            heapq.heappush(lower_half_max_heap, num * -1)
        elif difference <= -1:
            # lower_half_max_heap has more elements, should put this next number in upper_half_min_heap.
            heapq.heappush(upper_half_min_heap, num)
        else:
            # Can push to either, difference is 0.
            heapq.heappush(lower_half_max_heap, num * -1)