pygorithm/sort/heap_sort.py

47 lines
1.8 KiB
Python

"""
堆排序:将待排序的序列构建成一个小顶堆,此时,整个序列的最小值就是
堆顶根节点。将其与末尾元素进行交换,此时末尾就为最小值。这个最小值不再计算到堆内,
那么再将剩余的 n - 1 个元素重新构造成一个堆,这样会得到一个新的最小值。此时将该最
小值再次交换到新堆的末尾,这样就有了两个排序的值。重复这个过程,直到得到一个有序
序列。当然,小顶堆得到的是降序排序,大顶堆得到的才是升序排序。
"""
def sift_down(nums: list[int], n: int, i: int):
"""堆的长度为 n ,从节点 i 开始,从顶至底堆化"""
while True:
# 判断节点 i, l, r 中值最大的节点,记为 ma
l = 2 * i + 1 # i 的左子节点
r = 2 * i + 2 # i 的右子节点
ma = i
# 选择两个子节点中更大的那个
if l < n and nums[l] > nums[ma]:
ma = l
if r < n and nums[r] > nums[ma]:
ma = r
# 若节点 i 最大或索引 l, r 越界,则无须继续堆化,跳出
if ma == i:
break
# 交换两节点,将大节点向上移
nums[i], nums[ma] = nums[ma], nums[i]
# 循环向下堆化
i = ma
def heap_sort(nums: list[int]):
"""堆排序"""
# 建堆操作:堆化除叶节点以外的其他所有节点
for i in range(len(nums) // 2 - 1, -1, -1):
sift_down(nums, len(nums), i)
# 从堆中提取最大元素,循环 n-1 轮
for i in range(len(nums) - 1, 0, -1):
# 交换根节点与最右叶节点(交换首元素与尾元素)
nums[0], nums[i] = nums[i], nums[0]
# 以根节点为起点,从顶至底进行堆化
sift_down(nums, i, 0)
nums = [54, 32, 99, 18, 75, 31, ]
heap_sort(nums)
print(nums)