Python中的heapq模块源码详析

import heapq

data = [97, 38, 27, 50, 76, 65, 49, 13]
heap = []

for n in data:
 heapq.heappush(heap, n)

print('pop:', heapq.heappop(heap)) # pop: 13
print(heap) # [27, 50, 38, 97, 76, 65, 49]

import heapq

data = [97, 38, 27, 50, 76, 65, 49, 13]

heapq.heapify(data)

print('pop:', heapq.heappop(data)) # pop: 13
print(data) # [27, 38, 49, 50, 76, 65, 97]

def heapdown(heap, pos):
 endpos = len(heap)
 while pos < endpos:
 lchild = 2 * pos + 1
 rchild = 2 * pos + 2
 if lchild >= endpos: # 如果pos已经是叶节点，退出循环
  break
 childpos = lchild # 假设要交换的节点是左节点
 if rchild < endpos and heap[childpos] > heap[rchild]:
  childpos = rchild

 if heap[pos] < heap[childpos]: # 如果节点比子节点都小，退出循环
  break
 heap[pos], heap[childpos] = heap[childpos], heap[pos] # 交换
 pos = childpos

def heapup(heap, startpos, pos): # 如果是新增元素，startpos 传入 0
 while pos > startpos:
 parentpos = (pos - 1) // 2
 if heap[pos] < heap[parentpos]:
  heap[pos], heap[parentpos] = heap[parentpos], heap[pos]
  pos = parentpos
 else:
  break

for i in reversed(range(len(data) // 2)):
 heapdown(data, i)

def heappush(heap, item):
 """Push item onto heap, maintaining the heap invariant."""
 heap.append(item)
 _siftdown(heap, 0, len(heap)-1)

def _siftdown(heap, startpos, pos):
 newitem = heap[pos]
 # Follow the path to the root, moving parents down until finding a place
 # newitem fits.
 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):
 """Pop the smallest item off the heap, maintaining the heap invariant."""
 lastelt = heap.pop() # raises appropriate IndexError if heap is empty
 if heap:
  returnitem = heap[0]
  heap[0] = lastelt
  _siftup(heap, 0)
  return returnitem
 return lastelt

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 # 左节点，默认替换左节点
 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)

def heapify(x):
 """Transform list into a heap, in-place, in O(len(x)) time."""
 n = len(x)
 for i in reversed(range(n//2)):
  _siftup(x, i)