8.4. heapq — Heap queue algorithm

python标准库
8 数据类型
目标：堆是一个形如二叉树的数组，其父节点的值小于等于子节点的值，满足heap[k] <= heap[2k+1] and heap[k] <= heap[2k+2]，k从0开始，因此heap[0]是最小值，默认是最小堆。

使用场景：优先队列，任务调度，大磁盘排序，

可使用函数：

1. heapq.heappush(heap, item)
   将item插入堆中，维持堆的性质不变。
2. heapq.heappop(heap)
   取出最小值，第一个元素，并调整堆使之维持堆的性质不变。
3. heapq.heappushpop(heap, item)
   将item插入堆中，取出最小值，在维持一定长度的堆情况下很高效。总是返回heap[0]和item中较小的一个。
4. heapq.heapify(x)
   将列表x堆化，线性时间且就地。
5. heapq.heapreplace(heap, item)
   先返回最小值，然后插入item，在维持一定长度的堆情况下很高效。
6. heapq.merge(*iterables, key=None, reverse=False)
   合并多个已排序输入为一个排好序的输出，输出类型是迭代器。key是比较函数，reverse 默认升序。
7. heapq.nlargest(n, iterable, key=None)
   返回iterable提供的前n个最大的元素。key是比较函数。n较小时建议使用。
8. heapq.nsmallest(n, iterable, key=None)
   返回iterable提供的前n个最小的元素。key是比较函数。n较小时建议使用。

举例：

>>> def heapsort(iterable):
...     h = []
...     for value in iterable:
...         heappush(h, value)
...     return [heappop(h) for i in range(len(h))]
...
>>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

优先队列：
每个任务有自己的优先级，对任务排序。
问题：

1. 堆不是稳定排序，对两个相同优先级的不同名称任务，需要保持其插入顺序的有序。
2. 优先级可以比较，不同任务之间无法比较，会导致堆化中断。
3. 更改堆中某个任务的优先级。
4. 移除堆中某个任务，同时保持堆的性质。
   解决方案：
   1和2 使用计数器记住每个任务传进来的顺序，相同优先级的任务按照传入顺序排序。
   3和4都要先找到任务，使用字典记录任务对应其优先级和插入顺序。3的解决是将字典记录删除，重新插入新的优先级；原来堆中的那个任务并没有去掉，在pop时会产生疑惑。4的解决是使用删除标记（软删除）。

pq = []                         #  堆
entry_finder = {}               # 将任务映射于三元组[优先级，输入顺序，任务]
REMOVED = '<removed-task>'      # 删除标记
counter = itertools.count()     # 唯一的顺序计数器

def add_task(task, priority=0):
    '添加任务，更新任务优先级'
    if task in entry_finder:
        remove_task(task)
    count = next(counter)
    entry = [priority, count, task]
    entry_finder[task] = entry
    heappush(pq, entry)

def remove_task(task):
    '软删除任务.  Raise KeyError if not found.'
    entry = entry_finder.pop(task)
    entry[-1] = REMOVED

def pop_task():
    'Pop出最小元素. Raise KeyError if empty.'
    while pq:
        priority, count, task = heappop(pq)
        if task is not REMOVED:
            del entry_finder[task]
            return task
    raise KeyError('pop from an empty priority queue')

函数解析

def heappush(heap, item):
    """插入item，维持堆的性质"""
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1)

def _siftdown(heap, startpos, pos):
	# 自底向上循环比较置换的过程
    newitem = heap[pos]
    # 找到插入元素heap[pos]合适的位置。
    while pos > startpos:
        #找到父节点
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        #  位置变小：如果比父节点小，需要交换，向上继续走
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        # 就是当前位置
        break
    heap[pos] = newitem

def heappop(heap):
    """输出最小元素，维持堆的性质"""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
    # 将最后一个元素放在第一个位置，重新整理堆的第一个元素0
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup(heap, 0)
        return returnitem
    return lastelt

def heapreplace(heap, item):
    """先返回最小值，然后插入item，在维持一定长度的堆情况下很高效。"""
    returnitem = heap[0]    # raises appropriate IndexError if heap is empty
    heap[0] = item
    _siftup(heap, 0)
    return returnitem

def heappushpop(heap, item):
    """ 将item插入堆中，取出最小值，在维持一定长度的堆情况下很高效。总是返回heap[0]和item中较小的一个。如果item<heap[0]，则不需要插入"""
    if heap and heap[0] < item:
        item, heap[0] = heap[0], item
        _siftup(heap, 0)
    return item

def heapify(x):
    """列表x堆化in-place, in O(len(x)) time."""
    n = len(x)
    # 自底向上. 从最后一个非叶节点n//2-1开始。偶数情况下，最后节点只有左子树，奇数情况下，最后节点有左右子树。
    for i in reversed(range(n//2)):
        _siftup(x, i)

def _siftup(heap, pos):
    #  pos位置需要调整，将其一直向下调整到叶子结点，然后从叶子节点和传进的pos开始向上更新    _siftdown(heap, startpos, pos)
    endpos = len(heap)
    startpos = pos
    newitem = heap[pos]
    # 将孩子冒泡上去 如果孩子<父节点（当前需要更改的位置）
    childpos = 2*pos + 1    # 暂定左孩子为两个孩子的最小值
    while childpos < endpos:
        # 右孩子
        rightpos = childpos + 1
        # 右孩子存在，且右孩子<左孩子
        if rightpos < endpos and not heap[childpos] < heap[rightpos]:
            childpos = rightpos
        # 最小的孩子冒泡，当前节点向下走， 位于找到的最小的孩子的位置（位置变大）
        heap[pos] = heap[childpos]
        pos = childpos
        childpos = 2*pos + 1
    # 当前节点没有叶子节点，说明走到最下一层直接放元素
    heap[pos] = newitem
    # 从传进来的位置pos到最新固定的位置pos，将中间的父节点使其符合规则
    _siftdown(heap, startpos, pos)

def merge(*iterables, key=None, reverse=False):
    '''合并多个有序输入为一个有序输出，返回的是迭代器
    >>> list(merge([1,3,5,7], [0,2,4,8], [5,10,15,20], [], [25]))
        [0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25]
       使用key确定元素的排序
    >>> list(merge(['dog', 'horse'], ['cat', 'fish', 'kangaroo'], key=len))
        ['dog', 'cat', 'fish', 'horse', 'kangaroo']
    '''

    h = []
    h_append = h.append

    if reverse:
        _heapify = _heapify_max
        _heappop = _heappop_max
        _heapreplace = _heapreplace_max
        direction = -1
    else:
    # 升序排序
        _heapify = heapify
        _heappop = heappop
        _heapreplace = heapreplace
        direction = 1

    if key is None:
    # 没有指定比较函数
        for order, it in enumerate(map(iter, iterables)):
            try:
                next = it.__next__
                # 修改传进来的值，变为三元组[值, 传入序号，next方法]
                h_append([next(), order * direction, next])
            except StopIteration:
                pass
        # 将写好的三元组列表堆化
        _heapify(h)
        # 至少有两个才需要输出后重新堆化  
        while len(h) > 1:
            try:
                while True:
                	# 每次迭代输出最小的元素
                    value, order, next = s = h[0]
                    yield value
                    # 输出下一个元素之前，将原来的堆进行调整
                    s[0] = next()           # 耗尽时引发StopIteration
                    _heapreplace(h, s)      # 重新对第一个位置堆化
            except StopIteration:
                _heappop(h)                 # remove empty iterator
        if h:
            # 只有一个迭代器，输出value完成
            value, order, next = h[0]
            yield value
            yield from next.__self__
        return

    for order, it in enumerate(map(iter, iterables)):
        try:
            next = it.__next__
            value = next()
            # 修改传进来的值，变为四元组[需要比较的值 传入序号，值，next方法]
            h_append([key(value), order * direction, value, next])
        except StopIteration:
            pass
    _heapify(h)
    # 两个元素以上输出最小值后需要重新堆化
    while len(h) > 1:
        try:
            while True:
                # 每次迭代输出最小的元素
                key_value, order, value, next = s = h[0]
                yield value
                # 输出下一个元素之前，将原来的堆进行调整
                value = next()
                s[0] = key(value)
                s[2] = value
                # 重新对第一个位置堆化
                _heapreplace(h, s)
        except StopIteration:
            _heappop(h)
    if h:
        key_value, order, value, next = h[0]
        yield value
        yield from next.__self__

def nsmallest(n, iterable, key=None):
    """找到n个最小元素
       sorted(iterable, key=key)[:n]
    """

    #  n==1 等价于min()
    if n == 1:
        it = iter(iterable)
        sentinel = object()
        # sentinel是哨兵对象，如果it为空，返回的是哨兵对象。因为比较函数key不固定，因此传参可能是None.
        result = min(it, default=sentinel, key=key)
        return [] if result is sentinel else [result]

    #  n>=size, 直接使用sorted()
    try:
        size = len(iterable)
    # 当前实例没有属性函数len()
    except (TypeError, AttributeError):
        pass
    else:
        if n >= size:
            return sorted(iterable, key=key)[:n]

    # 没有指定key,使用最简单的装饰器
    if key is None:
        it = iter(iterable)
        # 元素转换为二元组[值，序列号]，首先取前n个最大堆化
        result = [(elem, i) for i, elem in zip(range(n), it)]
        # n为0时
        if not result:
            return result
        _heapify_max(result)
        
        top = result[0][0]
        order = n
        _heapreplace = _heapreplace_max
        for elem in it:
        # 找小于堆顶元素的，替换堆顶，并更新堆顶元素
            if elem < top:
                _heapreplace(result, (elem, order))
                top, _order = result[0]
                order += 1
        # 得到前n个最小，排序
        result.sort()
        return [elem for (elem, order) in result]

    # 更一般的方法
    it = iter(iterable)
    # 元素转换为三元组[需要比较的值，序列号，值]，首先取前n个最大堆化
    result = [(key(elem), i, elem) for i, elem in zip(range(n), it)]
    if not result:
        return result
    _heapify_max(result)
    
    top = result[0][0]
    order = n
    _heapreplace = _heapreplace_max
    for elem in it:
    # 找小于堆顶元素的，替换堆顶，并更新堆顶元素
        k = key(elem)
        if k < top:
            _heapreplace(result, (k, order, elem))
            top, _order, _elem = result[0]
            order += 1
    result.sort()
    return [elem for (k, order, elem) in result]



def nlargest(n, iterable, key=None):
    """Find the n largest elements in a dataset.
    Equivalent to:  sorted(iterable, key=key, reverse=True)[:n]
    """

    # Short-cut for n==1 is to use max()
    if n == 1:
        it = iter(iterable)
        sentinel = object()
        result = max(it, default=sentinel, key=key)
        return [] if result is sentinel else [result]

    # When n>=size, it's faster to use sorted()
    try:
        size = len(iterable)
    except (TypeError, AttributeError):
        pass
    else:
        if n >= size:
            return sorted(iterable, key=key, reverse=True)[:n]

    # When key is none, use simpler decoration
    if key is None:
        it = iter(iterable)
        result = [(elem, i) for i, elem in zip(range(0, -n, -1), it)]
        if not result:
            return result
        heapify(result)
        top = result[0][0]
        order = -n
        _heapreplace = heapreplace
        for elem in it:
            if top < elem:
                _heapreplace(result, (elem, order))
                top, _order = result[0]
                order -= 1
        result.sort(reverse=True)
        return [elem for (elem, order) in result]

    # General case, slowest method
    it = iter(iterable)
    result = [(key(elem), i, elem) for i, elem in zip(range(0, -n, -1), it)]
    if not result:
        return result
    heapify(result)
    top = result[0][0]
    order = -n
    _heapreplace = heapreplace
    for elem in it:
        k = key(elem)
        if top < k:
            _heapreplace(result, (k, order, elem))
            top, _order, _elem = result[0]
            order -= 1
    result.sort(reverse=True)
    return [elem for (k, order, elem) in result]