【Python】Python中的heapq模块详解,令人费解的_siftup与_siftdown函数
heapq模块基本介绍
heapq模块是python库中实现的小顶堆相关的函数的集合。对外暴露的接口函数有’heappush’, ‘heappop’, ‘heapify’, ‘heapreplace’, ‘merge’, ‘nlargest’, ‘nsmallest’, ‘heappushpop’。我挑选前两个进行详细的讲解。
小顶堆的定义和特性
在Python中,heapq模块提供了对小顶堆(Min Heap)的支持。小顶堆是一种数据结构,其中每个节点的值都小于或等于其子节点的值。这意味着堆的根节点始终是最小的元素。
堆的定义如下,n个关键字序列L[0…n-1】称为堆,当且仅当该序列满足∶
①L(i)>=L(2i+1)且L(i)>=L(2i+2)或
②L(i)<=L(2i+1)且L(i)<=L(2i+2)(0≤i≤n//2)
可以将该一维数组视为一棵完全二叉树,满足条件①的堆称为大根堆(大顶堆)。大根堆的最大元素存放在根结点。满足条件②的堆称为小根堆(小顶堆), 小根堆的定义刚好相反, 根结点是最小元素。
让我们通过一个简单的例子来理解小顶堆的特性:
import heapq
# 创建一个空的小顶堆
heap = []
# 往堆中添加元素
heapq.heappush(heap, 4)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)
# 弹出堆中最小的元素
min_element = heapq.heappop(heap)
print("堆中最小的元素是:", min_element)
print("剩余堆的元素:", heap)
堆中最小的元素是: 1
剩余堆的元素: [2, 4, 8]
上述例子中,我们使用heappush向堆中添加元素,然后使用heappop弹出最小的元素。小顶堆的特性确保每次弹出的元素都是堆中最小的。
heapq模块的作用
heapq模块主要用于处理堆数据结构,提供了一些函数接口来有效地进行堆操作。其中,heappush和heappop是两个常用的函数,让我们深入了解它们的作用。
heappush函数
heapq.heappush(heap, item)函数用于将元素item添加到堆heap中,并保持堆的特性。即在添加元素后,堆仍然是一个有效的小顶堆。
import heapq
heap = [1, 3, 5, 7, 9]
print("原始堆:", heap)
# 使用heappush添加元素到堆中
heapq.heappush(heap, 4)
print("添加元素后的堆:", heap)
原始堆: [1, 3, 5, 7, 9]
添加元素后的堆: [1, 3, 4, 7, 9, 5]
heappop函数
heapq.heappop(heap)函数用于弹出并返回堆heap中的最小元素。在弹出元素后,堆的特性仍然被维护。
import heapq
heap = [1, 3, 5, 7, 9]
print("原始堆:", heap)
# 使用heappop弹出最小元素
min_element = heapq.heappop(heap)
print("弹出的最小元素是:", min_element)
print("剩余堆的元素:", heap)
原始堆: [1, 3, 5, 7, 9]
弹出的最小元素是: 1
剩余堆的元素: [3, 7, 5, 9]
heappush函数与_siftdown函数
def heappush(heap, item):
"""Push item onto heap, maintaining the heap invariant."""
heap.append(item)
_siftdown(heap, 0, len(heap)-1)
# 'heap' is a heap at all indices >= startpos, except possibly for pos. pos
# is the index of a leaf with a possibly out-of-order value. Restore the
# heap invariant.
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 # 将新元素放置在正确的位置
当调用heapq.heappush(heap, item)时,heappush函数会将元素item添加到堆heap中,并确保堆的特性得以维持。在heappush函数内部,使用_siftdown函数来执行以下功能:将新添加的元素item放置在正确的位置,以保持小顶堆的性质。具体而言,它通过比较元素的值和其父节点的值,并在需要时进行交换,直到满足小顶堆的条件为止。因为heappush函数将新元素是放在堆底,然后让堆底元素上浮到合适的位置,所有称为新加元素的上浮操作,_siftdown的down有点让人混淆视听。用面向对象的编程方法改写heapq.heappush函数如下节。
改良后的heap_push函数
class HeapLittle:
def __init__(self):
self.size = 0
self.array = []
@staticmethod
def get_parent_index(child_index):
return (child_index - 1) // 2
def heap_push(self, item):
self.array.append(item)
self.size += 1
self.heap_item_up(0, self.size - 1)
def heap_item_up(self, start_position, end_position):
new_item = self.array[end_position]
iteration_index = end_position
iteration_parent_index = self.get_parent_index(iteration_index)
while iteration_parent_index >= start_position and new_item < self.array[iteration_parent_index]:
self.array[iteration_index] = self.array[iteration_parent_index]
iteration_index = iteration_parent_index
iteration_parent_index = self.get_parent_index(iteration_index)
self.array[iteration_index] = new_item
上面的heap_item_up函数改进了原版_siftdown函数的while循环的continue,break的逻辑,并增加函数封装更加表意。
heappop函数与_siftup函数
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 # 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位置的节点有子节点,选最小子节点的位置继续往下沉,直到沉到叶子节点。
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) #将沉到叶子节点(pos位置)往上浮到合适位置。
当调用heapq.heappop(heap)时,heappop函数会弹出并返回堆heap中的最小元素也就是堆顶元素,将堆的最后一个元素移动到堆顶,堆的长度减1。此时是新加入到堆顶的元素破坏了堆的性质,就需要把这个元素安置到合理位置。使用_siftup函数来执行以下功能:如果新加入到堆顶的元素有子节点,选它的最小子节点的位置往下沉,直到沉到叶子节点,此时叶子节点的位置假设为pos。那么就要执行_siftdown(heap, startpos, pos)函数将沉到的叶子节点(pos位置)往上浮到合适位置。以确保堆的特性得以维持。
分析上述_siftup函数,总是感觉怪怪的,是否有点脱裤子放屁–多次一举的操作。先将堆顶元素下沉到叶子节点,然后将叶子节点上浮到合适位置。这是否将操作量变大了?或许有我无法理解的深刻含义,读者如果有知道的可以告知一二。
改良后的heap_pop函数
def heap_pop(self):
"""Pop the smallest item off the heap, maintaining the heap invariant."""
last_item = self.array.pop() # raises appropriate IndexError if heap is empty
self.size -= 1
if self.size != 0:
return_item = self.array[0]
self.array[0] = last_item
# print(self)
# print(f"first not adjusted, item is {last_item}")
self.heap_item_down(0)
return return_item
return last_item
def heap_item_down(self, start_position):
new_item = self.array[start_position]
iteration_index = start_position
while self.has_left_child(iteration_index):
smaller_child_index = self.get_left_child_index(iteration_index)
right_child_index = self.get_right_child_index(iteration_index)
if self.has_right_child(iteration_index) and self.array[right_child_index] < self.array[smaller_child_index]:
smaller_child_index = right_child_index
if new_item > self.array[smaller_child_index]:
self.array[iteration_index] = self.array[smaller_child_index]
iteration_index = smaller_child_index
else:
break
# print(self)
# print(f"iteration_index is {iteration_index}, item value is {new_item}")
self.array[iteration_index] = new_item
# print(self)
# print("heap_item_down")
def has_left_child(self, index):
return self.get_left_child_index(index) < self.size
def get_left_child(self, index):
return self.array[self.get_left_child_index(index)]
def get_right_child(self, index):
return self.array[self.get_right_child_index(index)]
新加入到堆顶的元素破坏了堆的性质,就需要把这个元素安置到合理位置。使用heap_item_down函数来执行以下功能:如果新加入到堆顶的元素有子节点,并且它的子节点有比它值小的,那就选它的最小子节点的位置往下沉。直到已经沉到叶子节点了,或它的节点值都比它子节点值都小。
heapq模块与改进后的HeapLittle的验证
import numpy as np
import heapq
class HeapLittle:
def __init__(self):
self.size = 0
self.array = []
def print_heap(self):
for i in self.array:
print(i)
print("heap is above")
def __str__(self):
if self.size == 0:
return "Heap is empty"
else:
return str(self.array)
@staticmethod
def get_left_child_index(parent_index):
return 2 * parent_index + 1
@staticmethod
def get_right_child_index(parent_index):
return 2 * parent_index + 2
@staticmethod
def get_parent_index(child_index):
return (child_index - 1) // 2
def has_left_child(self, index):
return self.get_left_child_index(index) < self.size
def has_right_child(self, index):
return self.get_right_child_index(index) < self.size
def has_parent(self, index):
return self.get_parent_index(index) >= 0
def get_left_child(self, index):
return self.array[self.get_left_child_index(index)]
def get_right_child(self, index):
return self.array[self.get_right_child_index(index)]
def get_parent(self, index):
return self.array[self.get_parent_index(index)]
def heap_item_up(self, start_position, end_position):
new_item = self.array[end_position]
iteration_index = end_position
iteration_parent_index = self.get_parent_index(iteration_index)
while iteration_parent_index >= start_position and new_item < self.array[iteration_parent_index]:
self.array[iteration_index] = self.array[iteration_parent_index]
iteration_index = iteration_parent_index
iteration_parent_index = self.get_parent_index(iteration_index)
self.array[iteration_index] = new_item
def heap_item_down(self, start_position):
new_item = self.array[start_position]
iteration_index = start_position
while self.has_left_child(iteration_index):
smaller_child_index = self.get_left_child_index(iteration_index)
right_child_index = self.get_right_child_index(iteration_index)
if self.has_right_child(iteration_index) and self.array[right_child_index] < self.array[smaller_child_index]:
smaller_child_index = right_child_index
if new_item > self.array[smaller_child_index]:
self.array[iteration_index] = self.array[smaller_child_index]
iteration_index = smaller_child_index
else:
break
# print(self)
# print(f"iteration_index is {iteration_index}, item value is {new_item}")
self.array[iteration_index] = new_item
# print(self)
# print("heap_item_down")
def heap_push(self, item):
self.array.append(item)
self.size += 1
self.heap_item_up(0, self.size - 1)
def heap_pop(self):
"""Pop the smallest item off the heap, maintaining the heap invariant."""
last_item = self.array.pop() # raises appropriate IndexError if heap is empty
self.size -= 1
if self.size != 0:
return_item = self.array[0]
self.array[0] = last_item
# print(self)
# print(f"first not adjusted, item is {last_item}")
self.heap_item_down(0)
return return_item
return last_item
heap_little_object = HeapLittle()
array_normal = [np.random.randint(1, 50) for i in range(0, 10)]
print(array_normal)
for i in array_normal:
heap_little_object.heap_push(i)
print(heap_little_object)
for i in range(0, 10):
pop_item = heap_little_object.heap_pop()
print(heap_little_object)
print('python heapq')
array_heapq = []
for i in array_normal:
heapq.heappush(array_heapq, i)
print(array_heapq)
for i in range(0, 10):
pop_item = heapq.heappop(array_heapq)
print(array_heapq)
验证效果如下
[9, 13, 45, 3, 23, 2, 4, 49, 19, 36]
[9]
[9, 13]
[9, 13, 45]
[3, 9, 45, 13]
[3, 9, 45, 13, 23]
[2, 9, 3, 13, 23, 45]
[2, 9, 3, 13, 23, 45, 4]
[2, 9, 3, 13, 23, 45, 4, 49]
[2, 9, 3, 13, 23, 45, 4, 49, 19]
[2, 9, 3, 13, 23, 45, 4, 49, 19, 36]
[3, 9, 4, 13, 23, 45, 36, 49, 19]
[4, 9, 19, 13, 23, 45, 36, 49]
[9, 13, 19, 49, 23, 45, 36]
[13, 23, 19, 49, 36, 45]
[19, 23, 45, 49, 36]
[23, 36, 45, 49]
[36, 49, 45]
[45, 49]
[49]
Heap is empty
python heapq
[9]
[9, 13]
[9, 13, 45]
[3, 9, 45, 13]
[3, 9, 45, 13, 23]
[2, 9, 3, 13, 23, 45]
[2, 9, 3, 13, 23, 45, 4]
[2, 9, 3, 13, 23, 45, 4, 49]
[2, 9, 3, 13, 23, 45, 4, 49, 19]
[2, 9, 3, 13, 23, 45, 4, 49, 19, 36]
[3, 9, 4, 13, 23, 45, 36, 49, 19]
[4, 9, 19, 13, 23, 45, 36, 49]
[9, 13, 19, 49, 23, 45, 36]
[13, 23, 19, 49, 36, 45]
[19, 23, 45, 49, 36]
[23, 36, 45, 49]
[36, 49, 45]
[45, 49]
[49]
[]