迭代式快速排序Matplotlib动画报错修复求助

我想实现一个带Matplotlib动画演示的迭代式快速排序算法，用来展示排序进度。但用matplotlib.animation实现时，程序在首次排序并设置新基准点后停止，抛出错误：

for rect, val in zip(rects, A):
TypeError: zip argument #2 must support iteration

原排序函数代码 #

def partition(a, l, r):
    i = (l - 1)
    x = a[r]

    for j in range(l, r):
        if a[j] <= x:
            # increment index of smaller element
            i = i + 1
            a[i], a[j] = a[j], a[i]

    a[i + 1], a[r] = a[r], a[i + 1]
    return (i + 1)


def quickSortIterative(a, l, r):
    # Create an auxiliary stack
    size = r - l + 1
    stack = [0] * (size)

    # initialize top of stack
    top = -1

    # push initial values of l and h to stack
    top = top + 1
    stack[top] = l
    top = top + 1
    stack[top] = r

    # Keep popping from stack while is not empty
    while top >= 0:

        # Pop h and l
        r = stack[top]
        top = top - 1
        l = stack[top]
        top = top - 1

        # Set pivot element at its correct position in
        # sorted array
        if l >= r:
            return
        x = a[l]
        j = l
        for i in range(l + 1, r + 1):
            if a[i] <= x:
                j += 1
                a[j], a[i] = a[i], a[j]
            yield a
        a[l], a[j] = a[j], a[l]
        yield a
        # If there are elements on left side of pivot,
        # then push left side to stack
        for i in a:
            if i - 1 > l:
                top = top + 1
                stack[top] = l
                top = top + 1
                stack[top] = i - 1
                yield stack[top]

            # If there are elements on right side of pivot,
            # then push right side to stack

            if i + 1 < r:
                top = top + 1
                stack[top] = i + 1
                top = top + 1
                stack[top] = r
                yield stack[top]

原绘图代码 #

def draw2(n=None):
    n = int(n)
    a = [x + 1 for x in range(n)]
    random.shuffle(a)
    datasetName = 'Random'

    # generator object returned by the function
    generator = quickSortIterative(a, 0, n - 1)
    algoName = 'Quick Sort2'

    # style of the chart
    plt.style.use('fivethirtyeight')

    fig, ax = plt.subplots()

    # bar container
    bar_rects = ax.bar(range(len(a)), a, align="edge",
                       color='Blue')

    # setting the limits of x and y axes
    ax.set_xlim(0, len(a))
    ax.set_ylim(0, int(1.1 * len(a)))
    ax.set_title("ALGORITHM : " + algoName + "\n" + "DATA SET : " +
                 datasetName)

    text = ax.text(0.01, 0.95, "", transform=ax.transAxes, color="#E4365D")
    iteration2 = [0]

    def animate(A, rects, iteration2):
        for rect, val in zip(rects, A):
            # setting the size of each bar equal to the value of the elements
            rect.set_height(val)
        iteration2[0] += 1
        text.set_text("iterations : {}".format(iteration2[0]))

    # call animate function repeatedly
    anim = FuncAnimation(fig, func=animate,
                         fargs=(bar_rects, iteration2), frames=generator, interval=50,
                         repeat=False)
    plt.show()

错误原因 #

你的quickSortIterative生成器函数在后续步骤中yield了单个整数（栈的元素），而非排序过程中的数组a。FuncAnimation会将生成器的每个输出作为animate的第一个参数，当传入整数时，zip(rects, A)必然报错——因为整数不可迭代。

具体问题出在这段逻辑：

for i in a:
    if i - 1 > l:
        top = top + 1
        stack[top] = l
        top = top + 1
        stack[top] = i - 1
        yield stack[top]  # 此处yield单个整数，而非数组

    if i + 1 < r:
        top = top + 1
        stack[top] = i + 1
        top = top + 1
        stack[top] = r
        yield stack[top]  # 同样问题

这段代码混淆了数组元素和区间索引，且错误地yield栈元素，完全偏离了迭代式快排的逻辑。

修复方案 #

1. 将partition改为生成器，捕获每一次交换 #

为了展示分区过程中的每一步交换，让partition每次交换后yield数组副本（避免引用修改影响历史帧）：

def partition(a, l, r):
    i = l - 1
    pivot = a[r]

    for j in range(l, r):
        if a[j] <= pivot:
            i += 1
            a[i], a[j] = a[j], a[i]
            yield a.copy()  # 每次交换后yield当前数组状态

    a[i + 1], a[r] = a[r], a[i + 1]
    yield a.copy()
    return i + 1

2. 修正迭代式快排的栈逻辑 #

迭代式快排的栈应存储待处理的区间（左右边界），而非数组元素。修改quickSortIterative，遍历分区生成器的每一步输出：

def quickSortIterative(a, l, r):
    stack = []
    stack.append((l, r))

    while stack:
        current_l, current_r = stack.pop()
        if current_l >= current_r:
            continue
        
        # 遍历分区过程的每一步
        partition_gen = partition(a, current_l, current_r)
        try:
            while True:
                frame = next(partition_gen)
                yield frame
        except StopIteration as e:
            pivot_idx = e.value  # 获取分区后的基准索引
        
        # 将左右子区间压入栈（先压右区间，保证左区间优先处理）
        stack.append((current_l, pivot_idx - 1))
        stack.append((pivot_idx + 1, current_r))

3. 完整修复后的代码 #

整合所有修改后的可运行代码：

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random

def partition(a, l, r):
    i = l - 1
    pivot = a[r]

    for j in range(l, r):
        if a[j] <= pivot:
            i += 1
            a[i], a[j] = a[j], a[i]
            yield a.copy()

    a[i + 1], a[r] = a[r], a[i + 1]
    yield a.copy()
    return i + 1

def quickSortIterative(a, l, r):
    stack = []
    stack.append((l, r))

    while stack:
        current_l, current_r = stack.pop()
        if current_l >= current_r:
            continue
        
        partition_gen = partition(a, current_l, current_r)
        try:
            while True:
                frame = next(partition_gen)
                yield frame
        except StopIteration as e:
            pivot_idx = e.value
        
        stack.append((current_l, pivot_idx - 1))
        stack.append((pivot_idx + 1, current_r))

def draw2(n=None):
    n = int(n) if n else 10
    a = [x + 1 for x in range(n)]
    random.shuffle(a)
    datasetName = 'Random'

    generator = quickSortIterative(a, 0, n - 1)
    algoName = 'Iterative Quick Sort'

    plt.style.use('fivethirtyeight')
    fig, ax = plt.subplots()

    bar_rects = ax.bar(range(len(a)), a, align="edge", color='Blue')

    ax.set_xlim(0, len(a))
    ax.set_ylim(0, int(1.1 * len(a)))
    ax.set_title(f"ALGORITHM : {algoName}\nDATA SET : {datasetName}")

    text = ax.text(0.01, 0.95, "", transform=ax.transAxes, color="#E4365D")
    iteration2 = [0]

    def animate(A, rects, iteration2):
        for rect, val in zip(rects, A):
            rect.set_height(val)
        iteration2[0] += 1
        text.set_text(f"iterations : {iteration2[0]}")

    anim = FuncAnimation(fig, func=animate,
                         fargs=(bar_rects, iteration2), frames=generator, interval=50,
                         repeat=False)
    plt.show()

# 测试调用
draw2(10)

关键说明 #

• yield数组副本：直接yield原数组会导致后续修改覆盖历史帧，必须用a.copy()生成独立副本。
• 捕获生成器返回值：通过StopIteration的value属性获取分区后的基准索引，确保子区间正确入栈。
• 栈的正确使用：迭代式快排的栈存储待处理的区间边界，而非数组元素，这是原代码的核心逻辑错误。

内容的提问来源于stack exchange，提问作者User