使用Python+Matplotlib庫製作經典排序演算法可視化動畫

來自專欄 Bug工廠

近期在知乎上看到了幾個排序演算法可視化視頻，頓時興緻盎然，也想自己實現一次。最後在Matplotlib庫的幫助下達成了這個目的，做出了一條還算滿意的視頻，最後也開源了這個小玩具。

下面就是這個視頻成品，用Matplotlib庫自動生成原始素材，並用Premiere進行了後期處理。視頻中把九種排序演算法放在一起作為對比，分別是：

插入排序 希爾排序 選擇排序歸併排序 快速排序 堆排序冒泡排序 梳排序 猴子排序

先來一睹為快吧：

https://www.zhihu.com/video/991814541629505536

（加入猴子排序完全是想皮一下，畢竟一次成功的概率只有 ）

在構思方案的時候，我不可避免的走了一條由彎到直的路。當時的初始方案是，用python模擬各種排序演算法，然後逐幀生成步驟圖片，將圖片序列導入AE轉成視頻。後來覺得python直接寫圖片太麻煩，有沒有更方便的輪子？於是想到了Matplotlib庫可以繪製柱狀圖，剛好與我想要的圖片形式重合。後來改成了用Python+Matplotlib逐幀生成圖片，再將圖片序列導入AE轉成視頻。最後，又覺得圖片序列導入AE太麻煩，能不能直接輸出視頻？所幸Matplotlib既支持逐幀動畫，又可以和FFMpeg結合，直接以mp4格式將動畫輸出。excited！

至於如何得到排序演算法每個時刻的切片，我想過將幀編號和排序演算法的進度聯繫起來，邊播放邊獲取下一幀的數據。後來發現操作起來很有難度，完全可以先走一遍排序演算法，得到所有幀的數據，再逐幀播放。這樣雖然多耗了點內存，實現起來還是很簡單的。下面以基本的選擇排序為例介紹一下：

def selection_sort(data_set): ds = copy.deepcopy(data_set) for i in range(0, Data.data_count-1): for j in range(i+1, Data.data_count): if ds[j].value < ds[i].value: ds[i], ds[j] = ds[j], ds[i] return ds

選擇排序的代碼非常簡單，我們要做的就是在演算法比較有代表性的地方截取數據切片作為幀數據，然後同時處理幀數據，為某些重要的數值染色。最後的代碼是這樣的：

def selection_sort(data_set): # FRAME OPERATION BEGIN frames = [data_set] # FRAME OPERATION END ds = copy.deepcopy(data_set) for i in range(0, Data.data_count-1): for j in range(i+1, Data.data_count): # FRAME OPERATION BEGIN ds_r = copy.deepcopy(ds) frames.append(ds_r) ds_r[i].set_color(r) ds_r[j].set_color(k) # FRAME OPERATION END if ds[j].value < ds[i].value: ds[i], ds[j] = ds[j], ds[i] # FRAME OPERATION BEGIN frames.append(ds) return frames # FRAME OPERATION END

可以看到新代碼在原先的代碼上加了三塊處理幀數據的操作，並用注釋標了出來。第一塊：初始化幀列表，原始數據作為第一幀；第二塊：在第二層循環內部截取幀，並把第i個數據塗為紅色（r），第j個數據塗為黑色（k）；第三塊：在幀列表中加入已排序的數據作為最後一幀，並返回幀列表。

這樣，我們就可以用最直觀的方式看到選擇排序的過程以及i和j的意義：第i個元素一直在取j掃過部分的最小值。

當然，這只是這些排序演算法中較為簡單的截取及染色方案。還有一些演算法比較抽象，從排序過程中難以看出規律，比如堆排序。我給大根堆的每一層塗上了不同的顏色，並用紅色表示正在下沉或上浮的結點，用黑色表示紅色結點調整位置的過程中需要比較的孩子結點或父結點。這下排序過程總算直觀了些：

至於Matplotlib的繪圖和動畫部分，可以查閱官網文檔，這裡就不再贅述了。但是把動畫導出成mp4文件需要FFMpeg庫的幫助，具體做法這篇文章講得很好：matplotlib animation動畫保存(save函數)詳解。

我已經把這些代碼全部開源，並優化了一下用戶介面，有窗口播放動畫（9個演算法並列播放或單個演算法播放）、生成mp4和html的功能，可以參考README文檔使用。GitHub倉庫地址：

zamhown/sorting-visualizer?

github.com

說來慚愧，這是我寫README最認真的一次……歡迎star、fork或者提issue。如果想改進或者添加新的排序演算法歡迎提交pr。祝端午玩得愉快~

推薦閱讀：