深入堆（Heap）结构

class Heap:
    def __init__(self, key=lambda x:x):
        self.data = []
        self.K = key

    @staticmethod
    def _parent(idx):
        assert idx > 0
        return (idx-1) // 2  # same as >>1, not int()

    @staticmethod
    def _left(idx):
        return idx*2 + 1

    @staticmethod
    def _right(idx):
        return idx*2 + 2

    def heapify(self, data):
        self.data = data
        for i in reversed(range(len(data)//2)):
            self._trickle_down(i)

    def _trickle_down(self, idx=0):
        """ trickle down start at idx """
        while idx < len(self.data):
            left, right = Heap._left(idx), Heap._right(idx)
            tidx = idx
            if left<len(self.data) and self.K(self.data[tidx])<self.K(self.data[left]):
                tidx = left
            if right<len(self.data) and self.K(self.data[tidx])<self.K(self.data[right]):
                tidx = right
            if tidx == idx:
                return
            self.data[idx], self.data[tidx] = self.data[tidx], self.data[idx]
            idx = tidx

    def add(self, x):
        """ bubble up """
        idx = len(self.data)
        self.data.append(x)
        while idx > 0:
            parent = Heap._parent(idx)
            if self.K(self.data[idx]) > self.K(self.data[parent]):
                self.data[idx], self.data[parent] = self.data[parent], self.data[idx]
                idx = parent
                continue
            return

    def peek(self):
        ret = None
        if self.data:
            ret = self.data[0]
        return ret

    def pop(self):
        ret = None
        if self.data:
            ret = self.data[0]
            self.data[0] = self.data[-1]
            del self.data[-1]
            self._trickle_down()
        return ret

    def __bool__(self):
        return len(self.data) > 0

    def __len__(self):
        return len(self.data)

    def __repr__(self):
        return repr(self.data)


import random
h = Heap()
a = [i for i in range(16)]
random.shuffle(a)
for i in a:
    h.add(i)
print(h)
for i in range(len(a)+1):
    print(h.peek())
    assert h.peek() == h.pop()


h = Heap(lambda x:-x)
a = [i for i in range(16)]
random.shuffle(a)
for i in a:
    h.add(i)
print(h)
for i in range(len(a)+1):
    print(h.peek())
    assert h.peek() == h.pop()

h = Heap(lambda x:len(x))
a = [str(i)*i for i in range(9)]
random.shuffle(a)
for i in a:
    h.add(i)
b = [h.pop() for _ in range(len(h))]
print(b)

def heap_check(data, K=lambda x:x):
    if (size:=len(data)) == 0:
        return True
    idx = (size//2) - 1
    i2 = idx * 2
    if size & 1:
        t = data[i2+2] if K(data[i2+1])<K(data[i2+2]) else data[i2+1]
        if K(data[idx]) < K(t):
            return False
    else:
        if K(data[idx]) < K(data[i2+1]):
            return False
    for i in range(idx-1,-1,-1):
        i2 = i * 2
        t = data[i2+2] if K(data[i2+1])<K(data[i2+2]) else data[i2+1]
        if K(data[i]) < K(t):
            return False
    return True


import random
for n in range(2000):
    a = [i for i in range(n)]
    random.shuffle(a)
    h = Heap()
    h.heapify(a)
    assert len(h) == n
    assert heap_check(h.data)


K = lambda x: -x
for n in range(2000):
    a = [i for i in range(n)]
    random.shuffle(a)
    h = Heap(K)
    h.heapify(a)
    assert len(h) == n
    assert heap_check(h.data, K)

    def heapify(self, data):
        self.data = data
        for i in reversed(range(len(data)//2)):
            self._trickle_down(i)

# The child indices of heap index pos are already heaps, and we want to make
# a heap at index pos too.  We do this by bubbling the smaller child of
# pos up (and so on with that child's children, etc) until hitting a leaf,
# then using _siftdown to move the oddball originally at index pos into place.
#
# We *could* break out of the loop as soon as we find a pos where newitem <=
# both its children, but turns out that's not a good idea, and despite that
# many books write the algorithm that way.  During a heap pop, the last array
# element is sifted in, and that tends to be large, so that comparing it
# against values starting from the root usually doesn't pay (= usually doesn't
# get us out of the loop early).  See Knuth, Volume 3, where this is
# explained and quantified in an exercise.
#
# Cutting the # of comparisons is important, since these routines have no
# way to extract "the priority" from an array element, so that intelligence
# is likely to be hiding in custom comparison methods, or in array elements
# storing (priority, record) tuples.  Comparisons are thus potentially
# expensive.
#
# On random arrays of length 1000, making this change cut the number of
# comparisons made by heapify() a little, and those made by exhaustive
# heappop() a lot, in accord with theory.  Here are typical results from 3
# runs (3 just to demonstrate how small the variance is):
#
# Compares needed by heapify     Compares needed by 1000 heappops
# --------------------------     --------------------------------
# 1837 cut to 1663               14996 cut to 8680
# 1855 cut to 1659               14966 cut to 8678
# 1847 cut to 1660               15024 cut to 8703
#
# Building the heap by using heappush() 1000 times instead required
# 2198, 2148, and 2219 compares:  heapify() is more efficient, when
# you can use it.
#
# The total compares needed by list.sort() on the same lists were 8627,
# 8627, and 8632 (this should be compared to the sum of heapify() and
# heappop() compares):  list.sort() is (unsurprisingly!) more efficient
# for sorting.

def _siftup(heap, pos):
    endpos = len(heap)
    startpos = pos
    newitem = heap[pos]
    # Bubble up the smaller child until hitting a leaf.
    childpos = 2*pos + 1    # leftmost child position
    while childpos < endpos:
        # Set childpos to index of smaller child.
        rightpos = childpos + 1
        if rightpos < endpos and not heap[childpos] < heap[rightpos]:
            childpos = rightpos
        # Move the smaller child up.
        heap[pos] = heap[childpos]
        pos = childpos
        childpos = 2*pos + 1
    # The leaf at pos is empty now.  Put newitem there, and bubble it up
    # to its final resting place (by sifting its parents down).
    heap[pos] = newitem
    _siftdown(heap, startpos, pos)