更新matplotlib动画中当前矩形的颜色

Question

我想使用来自 matplotlib 的 FuncAnimation 动画冒泡排序算法。我的想法是生成一个随机的整数列表，并用这些值初始化一个条形图。然后我将 update_fig 函数与我的 bubble_sort 函数的生成器值一起应用。到目前为止一切顺利，我能够为整个冒泡排序算法制作动画，最后我的数字数组被排序。

但为了更好的可读性，我想更改当前矩形的颜色。例如：所有的矩形应该是蓝色的，当前的应该是红色的。

这是我的代码：

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

def update_fig(a, rects, iteration):
    """
    Function to update the plot
    :param a: the array with numbers
    :param rects: barrect for each number in the array
    :param iteration: the current iteration
    :return: None
    """
    for rect, val in zip(rects, a):
        rect.set_height(val)
        rect.set_color('grey')
    iteration[0] += 1
    text.set_text(f'# of operations {iteration[0]}')


def bubble_sort(a):
    """
    In-place Bubble Sort
    :param a: shuffled array
    :return: rearranged array
    """
    for i in range(len(a)):
        for j in range(0, len(a) - i - 1):
            if a[j] > a[j + 1]:
                tmp = a[j]
                a[j] = a[j + 1]
                a[j + 1] = tmp
            yield a

a = list(range(0,100))
random.shuffle(a)
generator = bubble_sort(a)
fig, ax = plt.subplots()
ax.set_title('Bubble Sort')
bar_rects = ax.bar(range(len(a)), a, align='edge')
ax.set_xlim(0, 100)
ax.set_ylim(0, int(1.07 * 100))

text = ax.text(0.02, 0.95, "", transform=ax.transAxes)
iteration = [0]
animate = animation.FuncAnimation(fig,
                                  func=update_fig,
                                  fargs=(bar_rects, iteration),
                                  frames=generator,
                                  interval=1,
                                  repeat=False,
                                  save_count=1500
                                  )
plt.show()

结果应该类似于 this。

我已经尝试在我的 update_fig 函数中设置 rect.set_color('orange') 但这会改变开始时绘图中所有矩形的颜色。

有什么想法吗？

Answer 1

一种方法是将当前值存储到全局变量中，可以在着色期间检查该变量。而且，当我们这样做的时候，也让迭代成为一个全局变量。

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

iteration = 0
current_val = None

def update_fig(a, rects):
     """
     Function to update the plot
     :param a: the array with numbers
     :param rects: barrect for each number in the array
     :param iteration: the current iteration
     :return: None
     """
     global iteration, current_val
     for rect, val in zip(rects, a):
          rect.set_height(val)
          rect.set_color('tomato' if val == current_val else 'grey')
     iteration += 1
     text.set_text(f'# of operations {iteration}')

def bubble_sort(a):
     """
     In-place Bubble Sort
     :param a: shuffled array
     :return: rearranged array
     """
     global current_val
     for i in range(len(a)):
          for j in range(0, len(a) - i - 1):
               if a[j] > a[j + 1]:
                    tmp = a[j]
                    a[j] = a[j + 1]
                    a[j + 1] = tmp
               current_val = a[j + 1]
               yield a

a = list(range(0, 100))
random.shuffle(a)
generator = bubble_sort(a)
fig, ax = plt.subplots()
ax.set_title('Bubble Sort')
bar_rects = ax.bar(range(len(a)), a, align='edge')
ax.set_xlim(0, 100)
ax.set_ylim(0, int(1.07 * 100))

text = ax.text(0.02, 0.95, "", transform=ax.transAxes)
animate = animation.FuncAnimation(fig,
                                  func=update_fig,
                                  fargs=(bar_rects,),
                                  frames=generator,
                                  interval=1,
                                  repeat=False,
                                  save_count=1500
                                  )
plt.show()

更新matplotlib动画中当前矩形的颜色

Update the color of current rectangle in matplot animation

python

matplotlib

matplotlib-animation