python实现排序算法过程的可视化

背景

今天老师居然布置个实验让我们做个排序算法的可视化出来，给了任务要求和参考代码链接

github链接为https://github.com/ZQPei/Sorting_Visualization
进去大致看了一下，算法应该都不难，能写出来。可视化它用的是pygame和cv。然后我也查了一下资料，发现我们python最熟悉的可视化库matplotlib里面的animation模块也能很好的实现动态可视化。对于这个模块的使用可以参考官方文档，另外这一篇博客介绍的也很有用
https://blog.csdn.net/sinat_36772813/article/details/77449122
万事俱备，开始撸代码

代码部分

首先我的代码结构很简单，就是一个sorting下存放排序算法和随机生成数据的代码，最后一个主文件调用这些就好。

main.py(我被要求写的是冒泡，希尔和堆）

import random
import sys
import re
from matplotlib import pyplot as plt
from matplotlib import animation
from sorting.data import Data
from sorting.bubblesort import bubble_sort
from sorting.shellsort import shell_sort
from sorting.heapsort import heap_sort

stype_dic = {'all': -1, 'shell-sort':0, 'heap-sort': 1,'bubble-sort': 2}
titles = [r'Shell Sort ($O(n \cdot log_2(n)^2)$)',r'Heap Sort ($O(n \cdot log_2(n))$)',r'Bubble Sort ($O(n^2)$)']
funs = [shell_sort,heap_sort,bubble_sort]

def create_original_data(dtype):
    data = []
    if dtype == 'random':
        data = list(range(1, Data.data_count + 1))
        random.shuffle(data)
    return data


def draw_chart(stype, original_data, frame_interval):
    fig = plt.figure(1, figsize=(16, 9))
    data_set = [Data(d) for d in original_data]
    axs = fig.add_subplot(111)
    axs.set_xticks([])
    axs.set_yticks([])
    plt.subplots_adjust(left=0.01, bottom=0.02, right=0.99, top=0.95,
                        wspace=0.05, hspace=0.15)
    frames = funs[stype](data_set)
    print('%s: %d frames.' % (re.findall(r'\w+ Sort', titles[stype])[0], len(frames)))

    def animate(fi):
        bars = []
        if (len(frames) > fi):
            axs.cla()
            axs.set_title(titles[stype])
            axs.set_xticks([])
            axs.set_yticks([])
            bars += axs.bar(list(range(Data.data_count)),  # X
                            [d.value for d in frames[fi]],  # data
                            1,  # width
                            color=[d.color for d in frames[fi]]  # color
                            ).get_children()
        return bars

    anim = animation.FuncAnimation(fig, animate, frames=len(frames), interval=frame_interval)
    return plt, anim


if __name__ == "__main__":
    try:
        Data.data_count = int(input('请输入想要排序数列的元素个数:(默认为10) '))
    except:
        Data.data_count = 10
    if len(sys.argv) > 1:
        stype = -1
        if len(sys.argv) > 2:
            if sys.argv[2] in stype_dic:
                stype = stype_dic[sys.argv[2]]
            else:
                print('Error: Wrong argument!')
                exit()
        stype_str = list(stype_dic.keys())[list(stype_dic.values()).index(stype)]
        dtype = 'random'
        if len(sys.argv) > 3:
            dtype = sys.argv[3]
            if dtype not in ('random'):
                print('Error: Wrong argument!')
                exit()
        od = create_original_data(dtype)

        if sys.argv[1] == 'play':
            try:
                fi = int(input('请设置帧间隔(默认为100)'))
            except:
                fi = 100
            plt, _ = draw_chart(stype, od, fi)
            plt.show()

再拿个堆排序的作为例子，另外两个就不贴了

import copy
from .data import Data

def heap_sort(data_set):
    frames = [data_set]

    ds = copy.deepcopy(data_set)

    for i in range(Data.data_count//2-1, -1, -1):
        heap_adjust(ds, i, Data.data_count, frames)
    for i in range(Data.data_count-1, 0, -1):

        ds[i], ds[0] = ds[0], ds[i]

        heap_adjust(ds, 0, i, frames)

    frames.append(ds)
    return frames


def heap_adjust(ds, head, tail, frames):

    i = head * 2 + 1
    while i < tail:

        if i + 1 < tail and ds[i].value < ds[i+1].value:
            i += 1

        ds_c = color(ds, tail)
        frames.append(ds_c)
        ds_c[i].set_color('k')
        ds_c[head].set_color('r')

        if ds[i].value <= ds[head].value:
            break
        ds[head], ds[i] = ds[i], ds[head]

        ds_c = copy.deepcopy(ds_c)
        frames.append(ds_c)
        ds_c[head], ds_c[i] = ds_c[i], ds_c[head]
        head = i
        i = i * 2 + 1

def color(ds, n):
    ds_c = copy.deepcopy(ds)
    head = 0
    tail = 1
    count = 1
    depth = 0
    colors = 'bmgcy'
    while head < n:
        for i in range(head, min(tail, n)):
            ds_c[i].set_color(colors[depth % len(colors)])
        head = tail
        count *= 2
        tail += count
        depth += 1
    return ds_c

data.py

class Data:
    data_count = 10
    def __init__(self, value):
        self.value = value
        self.set_color()

    def set_color(self, rgba = None):
        if not rgba:
            rgba = (0,
                    1 - self.value / (self.data_count * 2),
                    self.value / (self.data_count * 2) + 0.5,
                    1)
        self.color = rgba

成果

在命令行输入python main.py play +三种排序中的一种
比如python main.py play heap-sort
运行结果

好怕有同班同学搜这个实验搜到我