python写的堆排序动态演示程序

python写的堆排序动态演示程序

import matplotlib.pyplot as plt
import numpy as np
import time


def heapify(arr, n, i, ax):
    largest = i  # Initialize largest as root
    l = 2 * i + 1  # left = 2*i + 1
    r = 2 * i + 2  # right = 2*i + 2

    # See if left child of root exists and is greater than root
    if l < n and arr[i] < arr[l]:
        largest = l

    # See if right child of root exists and is greater than root
    if r < n and arr[largest] < arr[r]:
        largest = r

    # Change root, if needed
    if largest != i:
        arr[i], arr[largest] = arr[largest], arr[i]  # swap

        # 更新柱状图并暂停一段时间
        ax.clear()
        ax.bar(range(len(arr)), arr, color=['gray' if x != i and x != largest else 'red' for x in range(len(arr))])
        for j, num in enumerate(arr):
            ax.text(j, num + 1, str(num), horizontalalignment='center')
        plt.pause(1)

        # Heapify the root.
        heapify(arr, n, largest, ax)


def heap_sort(arr, ax):
    n = len(arr)

    # Build a maxheap.
    for i in range(n // 2 - 1, -1, -1):
        heapify(arr, n, i, ax)

    # Extract elements one by one
    for i in range(n - 1, 0, -1):
        arr[i], arr[0] = arr[0], arr[i]  # swap
        heapify(arr, i, 0, ax)

        # 更新柱状图并暂停一段时间
        ax.clear()
        ax.bar(range(len(arr)), arr, color='red')
        for j, num in enumerate(arr):
            ax.text(j, num + 1, str(num), horizontalalignment='center')
        plt.pause(1)


# 生成随机数组
arr = np.random.randint(1, 100, 12)

# 创建画布和轴对象
fig, ax = plt.subplots()
ax.set_title("Heap Sort")

# 绘制初始柱状图
ax.bar(range(len(arr)), arr)
for j, num in enumerate(arr):
    ax.text(j, num + 1, str(num), horizontalalignment='center')

# 刷新画布，暂停一段时间，然后开始堆排序
fig.canvas.draw()
plt.pause(1)

heap_sort(arr, ax)

# 更新柱状图并显示
ax.bar(range(len(arr)), arr)
for j, num in enumerate(arr):
    ax.text(j, num + 1, str(num), horizontalalignment='center')
fig.canvas.draw()
plt.show()