Sorting Algorithms are an important part of computer science. It’s great to be able to understand these. Being able to see the animations of these algorithms is a great way to understand them. I was recently asked to create a tutorial on how to create the animations on my Sorting Algorithms page.
To follow this guide, we’ll need to install the matplotlib library. We can do so with the terminal command below.
pip install matplotlibOverview of How to Create Sorting Animations in Python
In this post we’ll create a file that handles our animation function called animate_generator.py. To actually create the sorting animations, we’ll create a different file for each sorting algorithm.
In this post we’ll go over:
- Matplotlib Modules for Animation
- Introduction to matplotlib.animation.FuncAnimation
- Introduction to matplotlib.animation.PillowWriter
- Creating the Sorting Animation
- Setting Up the Figure to Animate
- Creating the Frame Animation Function
- Using FuncAnimation to Animate Function
- Save the GIF with the PillowWriter Class
- Full Code for Creating an Animation Function in Python
- Animating Different Sorting Algorithms
- Bubble Sort
- Heap Sort
- Insertion Sort
- Merge Sort
- Quick Sort
Matplotlib Modules for Animation
We’re going to use matplotlib, the de facto plotting library for Python. We’ve used matplotlib.pyplot in so many posts already such as how to plot a random dataset, an example of bifurcation, and how to create a recurrent neural network from scratch. This is the basic plotting library to plot functions and scatter plots in Python. For creating and saving an animation though, we’re going to use two classes from the matplotlib.animation library, FuncAnimation, and PillowWriter.
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriterIntroduction to matplotlib.animation.FuncAnimation
class matplotlib.animation.FuncAnimation(fig, func, frames=None, init_func=None, fargs=None, save_count=None, *, cache_frame_data=True, **kwargs)
FuncAnimation creates an animation with a passed in function. As we can see from the code signature above, FuncAnimation requires two parameters and takes up to seven optional parameters. We’ll use this to create our sorting animations. To learn more, read about how to create animations with FuncAnimation.
Introduction to matplotlib.animation.PillowWriter
class matplotlib.animation.PillowWriter(fps=5, metadata=None, codec=None, bitrate=None)
PillowWriter is the Python Image Library fork writer to save animations to a file. PillowWriter takes four parameters. None of the parameters are required, but we’ll be using the fps parameter. The other three are not commonly used. To learn more, read about how to save animations with PillowWriter in Python.
Creating the Sorting Animation Function
To create beautiful sorting animations in Python, the most important step is to actually create the function that will create our sorting animations. Our animate functions will take four mandatory parameters and one optional parameter. The four required parameters are _list, the list that we want to sort, generator, the sorting function in the form of a generator, title, the title of the file, save, whether or not we want to save the file, and save_frames, the number of frames to save.
We’ll begin our function by saving the length of the list to a variable, n, which we’ll use later to determine the axis size. Then we’ll get our figure and axis from subplots(), and set the title of the axis to the passed in title variable.
def animate(_list, generator, title, save, save_frames=1000):
n = len(_list)
fig, ax = plt.subplots()
ax.set_title(title)Setting Up the Figure to Animate
Now let’s set up the figure that we’re going to run our animations on. First thing we’ll need to do is set up the bars. We’re going to use a bar plot to create our sorting animation because we want to visualize the numbers moving around. Next, we’ll set the x and y limits. We set the xlim to the length of the list and the ylim to a little bit larger than the length of the list for a cleaner looking animation.
Then we’ll set some text to be displayed at the top corner of the image. Finally, we’ll set the iteration to a one element list with entry 0. Note that we’re not going to add elements to this list, just increment the element in it. Why do we use a list then? Because we don’t want to actually change the “value” of the variable itself, just the value that it holds. This is due to some under-the-hood implementation details of animation. Feel free to change it to an int from a list and play around with it yourself to see why.
bar_rects = ax.bar(range(len(_list)), _list, align="edge")
ax.set_xlim(0, n)
ax.set_ylim(0, int(1.1*n))
text = ax.text(0.01, 0.95, "", transform = ax.transAxes)
iteration = [0]Creating the Frame Animation Function
Now let’s create the function that we’ll pass into FuncAnimation to create the individual frames of our animation. We’ll define a function inside of our animation function called _animation (the underscore in front just denotes a private function/variable). This function will take three parameters, the array of numbers, the rects, representing the bars on the axis, and the iteration telling us which iteration it is.
In this function, we’ll set the values of all the bars in a frame, increment the first (and only) element of the iteration by 1, and then set the text to display which iteration we’re on.
def _animate(array, rects, iteration):
for rect, val in zip(rects, array):
rect.set_height(val)
iteration[0] += 1
text.set_text(f"iterations: {iteration[0]}")Using FuncAnimation to Animate Function
With everything set up, we can get to creating the actual animation. The matplotlib.animation.FuncAnimation class provides an easy way to stitch a bunch of frames together for an animation. All we need to do to create an animation is use FuncAnimation and pass it our figure, the _animate function we created earlier and the arguments we need for it through fargs, and the generator function for each frame. We’ll also pass an interval between frames, whether or not we want to repeat the animation (I set this to False, but this is actually trivial), and the number of frames we want to save in the case that we save this animation.
anim = FuncAnimation(fig, func=_animate,
fargs=(bar_rects, iteration), frames=generator, interval=10,
repeat=False, save_count=save_frames)Save the GIF with the PillowWriter Function
Finally, we’ll use the matplotlib.animation.PillowWriter class to create a writer to save our animation with. Remember that we had save as a parameter for our animate function to dictate whether or not we’ll actually save the outputted animation as a file. If we choose to save the file, we’ll use PillowWriter to create a writer with an fps of 30. Then we’ll call the animation object to save it to a .gif file using the PillowWriter object.
Finally, regardless of whether or not we save the .gif file, we’ll show the figure so we can see the animated sorting algorithm.
if save:
writer = PillowWriter(fps=30)
anim.save(title + ".gif", writer=writer)
plt.show()Full Code for Creating an Animation Function in Python
Here’s the full code for creating an animation function we can use to animate our sorting algorithms and save them as .gif files.
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
def animate(_list, generator, title, save, save_frames=1000):
n = len(_list)
fig, ax = plt.subplots()
ax.set_title(title)
bar_rects = ax.bar(range(len(_list)), _list, align="edge")
ax.set_xlim(0, n)
ax.set_ylim(0, int(1.1*n))
text = ax.text(0.01, 0.95, "", transform = ax.transAxes)
iteration = [0]
def _animate(array, rects, iteration):
for rect, val in zip(rects, array):
rect.set_height(val)
iteration[0] += 1
text.set_text(f"iterations: {iteration[0]}")
anim = FuncAnimation(fig, func=_animate,
fargs=(bar_rects, iteration), frames=generator, interval=10,
repeat=False, save_count=save_frames)
if save:
writer = PillowWriter(fps=30)
anim.save(title + ".gif", writer=writer)
plt.show() Animating Different Sorting Functions
Now that we’ve created the sort algorithm animation function, we’re ready to apply it. Let’s apply it to five common sorting algorithms we covered in the Sorting Algorithms You Need to Know post. We’ll animate Bubble Sort, Heap Sort, Insertion Sort, Merge Sort, and Quick Sort.
Bubble Sort Animation with Python
The first sorting algorithm we’ll animate with our animate function is bubble sort. Bubble sort is one of the first sorting algorithms taught in a beginner data structures and algorithms course on computer science.
The steps for bubble sort can simply be expressed like so:
- Use a
forloop to loop through the list- For each element in the list, loop through each element before it
- If the element before the element we’re checking is less than the element after it, swap the elements
- For each element in the list, loop through each element before it
This results in a “bubbling” of the largest elements to the top (end) of the list.
In the example code we begin by importing the animate function from animate_generator.py and the random module. Our bubblesort function will take one parameter, an array. The first thing we’ll do is assign the length of the array to a variable for simplicity. Then we’ll initialize a double for loop. The outer loop will loop through each element, the inner loop will loop through each element less than the current index of the outer loop.
Inside the inner loop, we’ll compare each element to the one after it. If the element is greater than the element after it, we’ll swap them. Then we’ll yield the changed array. At the end of the double for loop, our function will give us a generator function that will contain a lazy iterator “containing” each swap in the bubble sort algorithm.
Finally, we’ll create an array of length 50, shuffle it, and call the animate function on it to create our animation.
from animate_generator import animate
import random
def bubblesort(arr):
n = len(arr)
for i in range(n-1):
for j in range(n-i-1):
if arr[j] > arr[j+1]:
arr[j], arr[j+1] = arr[j+1], arr[j]
yield arr
array = [i for i in range(50)]
random.shuffle(array)
animate(array, bubblesort(array), "bubble sort", True)Heap Sort Animation with Python
This is my favorite sorting animation. I think heap sort is the most beautiful sort. The next sorting algorithm we’ll animate is Heap Sort. Heap sort is more complex than bubble sort. It’s a recursive sorting animation that “heaps” the list over and over again and then sorts the heap.
from animate_generator import animate
import randomHeap Function for Heap Sort
We need to create a heap function to turn the list into heaps for sorting. Our heap function will take three parameters, arr, the array, n, the length of the array, and i, the root element we’re heaping from. Inside the function we assign the left and right children of the root node as the first and second indices after doubling the root index.
Then we compare to see if the left child’s index is in the array and if the array entry at that index is greater than the ith index. If it is, we assign the left child’s index as the new root index. If the right child’s index is less than the size of the array and if the entry at the current root (which could be the left child or the original i) is less than the value of the right child, we assign the right child’s index to the root.
If we change the value of the root index during either of those checks, we’ll swap the values at the new root and the original i. Then we’ll recursively call the heap function again with the new root value as the new i.
def heap(arr, n, i):
root = i
lc = 2*i + 1
rc = 2*i + 2
if lc < n and arr[i] < arr[lc]:
root = lc
if rc < n and arr[root] < arr[rc]:
root = rc
if root != i:
arr[i], arr[root] = arr[root], arr[i]
heap(arr, n, root)Heap Sort Function
Now that we’ve created the heap function for heap sort, let’s create the heap_sort function that will execute the heap function on the whole array. The heap_sort function takes one parameter, an unsorted list. The first thing we do in our heap_sort function is store the length of the array as a variable.
For each of the elements in the first half of the array, starting from the middle and working backwards, we call the heap function with that element as the root. We also yield each of the resulting arrays for each heap call.
Then we’ll start from the end and work our way backwards again to the starting index. For each one, we’ll swap the focus element, i, and the first element, and yield the array. Then we call the heap function on the array with the first element as the root element and yield that array.
Finally, we’ll shuffle an array of size 100 and call the animate function on it to create the animation. We can sort 100 here instead of just 50 because this sort is WAY faster than bubble sort. Read the post about Sorting Functions to learn more about the time complexity comparison.
def heap_sort(arr):
n = len(arr)
for i in range(n//2, -1, -1):
heap(arr, n, i)
yield arr
for i in range(n-1, 0, -1):
arr[i], arr[0] = arr[0], arr[i]
yield arr
heap(arr, i, 0)
yield arr
array = [i for i in range(100)]
random.shuffle(array)
animate(array, heap_sort(array), "heap sort", False)Full Code for Heap Sort Algorithm with Animation in Python
from animate_generator import animate
import random
def heap(arr, n, i):
root = i
lc = 2*i + 1
rc = 2*i + 2
if lc < n and arr[i] < arr[lc]:
root = lc
if rc < n and arr[root] < arr[rc]:
root = rc
if root != i:
arr[i], arr[root] = arr[root], arr[i]
heap(arr, n, root)
def heap_sort(arr):
n = len(arr)
for i in range(n//2, -1, -1):
heap(arr, n, i)
yield arr
for i in range(n-1, 0, -1):
arr[i], arr[0] = arr[0], arr[i]
yield arr
heap(arr, i, 0)
yield arr
array = [i for i in range(100)]
random.shuffle(array)
animate(array, heap_sort(array), "heap sort", False)Insertion Sort Animation with Python
Insertion sort is called such because of the way it creates the sorted list. The insertion sort algorithm inserts sorted elements into a list as it grows. Each time it encounters an element, it inserts it into the already sorted list of elements before it.
Our insertion_sort algorithm will take one parameter, the unsorted list. We’ll loop through each element in the list and check the next element. While the next element is less than the focus element and we’re not at the 0 index, we’ll swap the element with the one before it and yield the changed array. Then we’ll set the i+1th element to the next element that we initially got at the beginning of our loop and yield this updated array.
Finally, to make the animation, we’ll create a randomized list of 50 elements and pass it to the animate function from animation_generator.py.
from animate_generator import animate
import random
def insertion_sort(arr):
for i in range(len(arr)-1):
next = arr[i+1]
while arr[i] > next and i >= 0:
arr[i+1] = arr[i]
i -= 1
yield arr
arr[i+1] = next
yield arr
array = [i for i in range(50)]
random.shuffle(array)
animate(array, insertion_sort(array), "insertion sort", False)Merge Sort Animation with Python
The merge sort algorithm is a sorting algorithm that continuously breaks down an unsorted list in half until it gets to just 2 elements. Once it gets that list, it sorts those elements and merges them until we get a wholly sorted list. The only two imports we’ll have to import are the animate function from animate_generator and the random library.
from animate_generator import animate
import randomMerge Function for Merge Sort
First let’s create a merge function, this function will take four parameters, the list we want to merge sort, the start index, the middle index, and the end index. The first thing we’ll do is create an empty list of results and then set the left and right indices. The left index will be the start index and the right index will be one past the middle index. These indicate where we’ll start organizing the left and right side of the list.
The next thing we’ll do is append to our result array. We will go through the left and right side of the array we passed in. If the element in the left index is less than the element in the right index, we append the value from the left index into the result array and increment the left index. Otherwise we’ll append the value from the right index into the result array and increment the right index.
After we reach the end of either the left or right side, we go through the rest of the left array and right array (only one of these will still have terms actually) and append the rest of that list to the result list. After each of these, we’ll loop through the result list and set the values in the original to the right values and yield the updated list.
def merge(_arr, start, mid, end):
result = []
left_index = start
right_index = mid + 1
while left_index <= mid and right_index <= end:
if _arr[left_index] < _arr[right_index]:
result.append(_arr[left_index])
left_index+=1
else:
result.append(_arr[right_index])
right_index+=1
while left_index <= mid:
result.append(_arr[left_index])
left_index+=1
while right_index <= end:
result.append(_arr[right_index])
right_index+=1
for index, value in enumerate(result):
_arr[start+index] = value
yield _arrSort Function for Merge Sort
As I said earlier, the merge sort function takes two functions to do, a merge function and a sort function. The sort function will take three parameters, the unsorted list, the start index, and the end index. The first thing we’ll do in the sort function is check for validity of the start and end indices and then create a middle index.
From there, all we need to do is recursively sort the left and right sides and then merge the two. We’ll yield from each of those calls to get the intermediate steps and finally we’ll yield the sorted array. To see our animation in action, we simply create a randomized list of 100 integers and pass it to the animate function.
def sort(_arr, start, end):
if end <= start:
return
mid = start + ((end-start+1)//2)
yield from sort(_arr, start, mid-1)
yield from sort(_arr, mid, end)
yield from merge(_arr, start, mid-1, end)
yield _arr
N=100
array = [i for i in range(N)]
random.shuffle(array)
animate(array, sort(array, 0, N-1), "merge sort", True)Full Code for Merge Sort Algorithm with Animation in Python
from animate_generator import animate
import random
def merge(_arr, start, mid, end):
result = []
left_index = start
right_index = mid + 1
while left_index <= mid and right_index <= end:
if _arr[left_index] < _arr[right_index]:
result.append(_arr[left_index])
left_index+=1
else:
result.append(_arr[right_index])
right_index+=1
while left_index <= mid:
result.append(_arr[left_index])
left_index+=1
while right_index <= end:
result.append(_arr[right_index])
right_index+=1
for index, value in enumerate(result):
_arr[start+index] = value
yield _arr
def sort(_arr, start, end):
if end <= start:
return
mid = start + ((end-start+1)//2)
yield from sort(_arr, start, mid-1)
yield from sort(_arr, mid, end)
yield from merge(_arr, start, mid-1, end)
yield _arr
N=100
array = [i for i in range(N)]
random.shuffle(array)
animate(array, sort(array, 0, N-1), "merge sort", True)Quicksort Animation with Python
The last sorting algorithm we’ll go over is the quicksort algorithm. The quicksort algorithm is based around choosing a “pivot” and sorting from start to end around that pivot by sorting the numbers less than it, then the numbers greater than it.
As always, we need to import the animate function and the random library. From there we’ll make a quicksort function that takes three parameters, the unsorted list, the start index, and the end index. Our function will start off by validating the indices and then defining a pivot and a swap index (as the start index for this example, but it’s actually up to you how you want to define it).
For each of the elements after the start index and before the end index, we’ll compare the element against the pivot. If the element is less than or equal to the pivot (it will never be equal in our example), we’ll increment the swap index and swap the ith element and the element in the swap index. We’ll yield the resulting array for each of these loops.
After we’re done traversing from the start index up, we swap the elements in the swap and start indices and yield the updated list. Then we yield from a quicksort of the new start index to one under the current swap index and a quicksort of one over the current swap index to the end index. Note that we don’t need to include the swap index from either of the recursive quicksort calls because it’s already in place.
from animate_generator import animate
import random
def quicksort(array, start, end):
if start >= end:
return
pivot = array[start]
swap = start
for i in range(start+1, end+1):
if array[i] <= pivot:
swap += 1
array[swap], array[i] = array[i], array[swap]
yield array
array[start], array[swap], = array[swap], array[start]
yield array
yield from quicksort(array, start, swap-1)
yield from quicksort(array, swap+1, end)
array = [i for i in range(100)]
random.shuffle(array)
animate(array, quicksort(array, 0, 99), "quicksort", False)I run this site to help you and others like you find cool projects and practice software skills. If this is helpful for you and you enjoy your ad free site, please help fund this site by donating below! If you can’t donate right now, please think of us next time.
4 replies on “Create Beautiful Sorting Animations in Python”
[…] structures like stacks, queues, and dequeues, linked lists and binary trees, and algorithms like sorting algorithms, tree algorithms, and Dijsktra’s algorithm. Now, let’s take a look at another important graph […]
[…] structures like stacks, queues, and dequeues, linked lists and binary trees, and algorithms like sorting algorithms, tree algorithms, and Dijsktra’s algorithm. Now, let’s take a look at another important graph […]
[…] Create Beautiful Sorting Animations with Python […]
[…] Previous Post Previous post: Create Animations with FuncAnimation in PythonNext Post Next post: Create Beautiful Sorting Animations in Python […]