# Heapsort with Numba and Cython

**Updated** April 20, 2022 following usefull comments by @scoder. Thank you very much for your pull requests!

*Heapsort* is a classical sorting algorithm. We are going into a little bit of theory about the algorithm, but refer to Corman et al. [1] for more details, or the heapsort wikipedia page.

In this post, we are going to implement the classical *heapsort* in Python, Python/Numba and Cython. The regular implementation is array-based and performed in-place. We use 0-based indices. Note that this is not a stable sorting method [keeping items with the same key in the original order].

## Imports

```
from itertools import cycle
from binarytree import build
import cython
from numba import njit
import numpy as np
import perfplot
%load_ext cython
SD = 124 # random seed
rng = np.random.default_rng(SD) # random number generator
```

Language/package versions:

```
Python implementation: CPython
Python version : 3.9.12
IPython version : 8.2.0
binarytree : 6.5.1
matplotlib : 3.5.1
perfplot : 0.10.2
numpy : 1.21.5
cython : 0.29.28
numba : 0.55.1
```

## Float array creation

We assume that we want to sort an NumPy array of 64-bit floating-point numbers :

```
n = 10
A = rng.random(n, dtype=np.float64)
A
```

```
array([0.78525311, 0.78585936, 0.96913602, 0.74805977, 0.65555081,
0.93888454, 0.17861445, 0.58864721, 0.44279917, 0.34884712])
```

## Binary trees

*Heapsort* is based on binary heap data structure, which relies on a nearly complete binary tree:

a [nearly] complete binary tree is a binary tree in which every level, except possibly the last, is completely filled, and all nodes in the last level are as far left as possible.

In a binary heap, the elements in an array are mapped to the tree nodes of a virtual binary tree. Assuming that we know the maximum number of elements in the heap, the array representation of the tree is more convenient than a linked structure. Note that it is also possible to use and array that can grow or shrink dynamically when threshold sizes are reached. Here is a figure showing the array elements and the corresponding tree nodes:

So we can see the array as a representation of the tree (implicit data structure): the node at the top [`A[0]`

] of the tree is called the **root**, and the ones at the bottom, without children, are called the **leaves** [`A[5]`

, `A[6]`

, `A[7]`

, `A[8]`

, `A[9]`

]. An advantage of binary trees is that it is easy to navigate among the nodes. Tree node indexing goes from the root to the leaves and from left to right. Level `i`

, the parent node index can be easily found as `(i - 1) // 2`

:

```
def parent(i):
assert i > 0
return (i - 1) // 2
```

`parent(1)`

```
0
```

`parent(8)`

```
3
```

The left child has index `2 * i + 1`

and the right child `2 * (i + 1)`

:

```
def left_child(i):
return 2 * i + 1
def right_child(i):
return 2 * (i + 1)
```

`left_child(0)`

```
1
```

`right_child(0)`

```
2
```

The **depth** of a node is the number of edges from the root node to the target node. The **height**

So that we have

## Binary heap

Here is the definition of a heap [2]:

A heap is a set of nodes with keys arranged in a complete heap-ordered binary tree, represented as an array.

A binary heap is binary tree data structure that satisfies a heap property. There are two kinds of heaps : min- and max-heaps satisfying respectively a min-heap or a max-heap property. In our case, we are going to use a max-heap to easily implement the *heapsort* algorithm. A min-heap would be used if sorting the array in a descending order, or would imply extra space/work to sort in ascending order.

**max-heap property** : A value of a given node `i`

is not larger than the value of its parent node `parent(i)`

Thus value of the root node `A[0]`

of a max-heap is greater than or equal to all the tree values. This is also true for any sub-tree of the binary heap: the root node of any sub-tree `A[i]`

is greater than or equal to all the values in this sub-tree.

An important point is that not all the array elements might be in the heap. We differentiate the heap `size`

from the array length `size`

elements in the left part of the array: `A[0:size]`

[with a Python slicing indexing]. All remaining elements `A[size:n]`

are not in the heap, implicitely.

Initially, given the above array `A`

, it mostly likely does not satisfy the max-heap property, which would imply in our case:
`A[0] >= A[1], A[0] >= A[2], A[1] >= A[3], A[1] >= A[4], A[3] >= A[7], A[3] >= A[8 ], A[2] >= A[5], A[2] >= A[6]`

. Indeed we have:

`A[0] >= A[1]`

```
False
```

`A[0] >= A[2]`

```
False
```

We need to build the max-heap from an array by moving around the elements. In order to do that, we use a function called `max_heapify`

from the leaves to the root of the tree, in a bottom-up manner. So let's implement this function.

## Max_heapify

The `max_heapify(A, size, i)`

function presupposes that sub-trees rooted at the children nodes of `i`

do satisfy the max-heap property. But the `i`

-th node may violate it, i.e. `A[i] < A[left_child(i)]`

or `A[i] < A[right_child(i)]`

, assuming that the children are in the heap. If this is the case, the `i`

-th node is swapped with its child with `largest`

value:

`A[i], A[largest] = A[largest], A[i]`

This process is then repeated from the largest children node. Eventually, the sub-tree rooted at `i`

satisfies the max-heap property after a call to `max_heapify(A, size, i)`

.

This process is also refered to as **sift down**: move a value violating the max-heap property down the tree by exchanging it with the largest of its 2 children values. Now it would also be possible to use the exact opposite process: **sift up**, which would start from a leaf node, and move its value up the tree by exchanging it with the parent node value, if violating the max-heap property. However, building a max-heap with the sift up process has a larger complexity than with sift down. I found this explanation by @alestanis on Stack Overflow really clear (The question was **Why siftDown is better than siftUp in heapify?**):

So here is a Python implementation:

```
def max_heapify(A, size, node_idx=0):
largest = node_idx
l = left_child(largest)
r = right_child(largest)
if (l < size) and (A[l] > A[largest]):
largest = l
if (r < size) and (A[r] > A[largest]):
largest = r
if largest != node_idx:
A[node_idx], A[largest] = A[largest], A[node_idx] # exchange 2 nodes
max_heapify(A, size, largest)
```

Note that this function is recursive. Let's have a look at the tree values of our example:

```
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
___________________0.78525___________
/ \
__________0.78586___________ ___0.96914___
/ \ / \
___0.74806__ ___0.65555 0.93888 0.17861
/ \ /
0.58865 0.4428 0.34885
```

This array is random, but we can observe that the max-heap property is satisfied everywhere except at the root! The sub-trees rooted at nodes 1 and 2 do happen to satisfy the max-heap property already. So let's call the `max_heapify`

function on the root node.

```
max_heapify(A, len(A), node_idx=0)
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
___________________0.96914___________
/ \
__________0.78586___________ ___0.93888___
/ \ / \
___0.74806__ ___0.65555 0.78525 0.17861
/ \ /
0.58865 0.4428 0.34885
```

The root node `A[0] = 0.78525`

moved from node index 0 to 2, and then from 2 to 5. Nodes 2 and 5 moved one step up in the process.

Steps:

- node 0 is compared with nodes 1 and 2. Node 0 is smaller than one of its children [0.78525 < 0.96914]. Node 2 is the child with largest value : exchange node 0 and 2.
- node 2 is compared with nodes 5 and 6. Node 0 is smaller than one of its children [0.78525 < 0.93888]. Node 5 is the child with largest value : exchange node 2 and 5.
- Node 5 is a leaf, so stop.

Let's generate a new random array to see another example:

```
A = rng.random(n, dtype=np.float64)
A
```

```
array([0.3309295 , 0.15936868, 0.98946349, 0.25711078, 0.71576487,
0.50588512, 0.66411132, 0.70234247, 0.05208023, 0.06009649])
```

```
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
__________________0.33093___________
/ \
___________0.15937__________ ___0.98946___
/ \ / \
___0.25711___ ___0.71576 0.50589 0.66411
/ \ /
0.70234 0.05208 0.0601
```

If we start from the leaves toward the root, we can notice a max-heap violation at node 3. So we call `max_heapify`

on that node:

```
max_heapify(A, len(A), node_idx=3)
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
__________________0.33093___________
/ \
___________0.15937__________ ___0.98946___
/ \ / \
___0.70234___ ___0.71576 0.50589 0.66411
/ \ /
0.25711 0.05208 0.0601
```

Now let's take care of node 1:

```
max_heapify(A, len(A), node_idx=1)
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
__________________0.33093___________
/ \
___________0.71576__________ ___0.98946___
/ \ / \
___0.70234___ ___0.15937 0.50589 0.66411
/ \ /
0.25711 0.05208 0.0601
```

And finally we need to fix the root value:

```
max_heapify(A, len(A), node_idx=0)
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
__________________0.98946___________
/ \
___________0.71576__________ ___0.66411___
/ \ / \
___0.70234___ ___0.15937 0.50589 0.33093
/ \ /
0.25711 0.05208 0.0601
```

We actually built a max-heap manually. In the following let's see the general method to build it.

## Build_max_heap

The method used to build the max-heap with a sift-down-based heapify is called Floyd's heap construction:

Floyd's algorithm starts with the leaves, observing that they are trivial but valid heaps by themselves, and then adds parents. Starting with element n/2 and working backwards, each internal node is made the root of a valid heap by sifting down. The last step is sifting down the first element, after which the entire array obeys the heap property.

This bottom-up heap construction technique runs in

```
def build_max_heap(A):
size = len(A)
node_idx = size // 2 - 1 # last non-leaf node index
for i in range(node_idx, -1, -1):
max_heapify(A, size, node_idx=i)
return size
```

```
A = rng.random(n, dtype=np.float64)
A
```

```
array([0.94535273, 0.25035043, 0.40579395, 0.27596342, 0.30065296,
0.36667218, 0.14878984, 0.34834079, 0.59713033, 0.99416357])
```

```
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
___________________0.94535___________
/ \
___________0.25035___________ ___0.40579___
/ \ / \
___0.27596___ ___0.30065 0.36667 0.14879
/ \ /
0.34834 0.59713 0.99416
```

```
_ = build_max_heap(A)
tree_values = build(list(A.round(5)))
print(tree_values)
```

```
___________________0.99416___________
/ \
___________0.94535___________ ___0.40579___
/ \ / \
___0.59713___ ___0.30065 0.36667 0.14879
/ \ /
0.34834 0.27596 0.25035
```

## Heapsort

The classical *heapsort* has two steps:

1 - build the heap

2 - destroy the heap by removing the root from the heap and moving it to the end of the heap

Step 2 corresponds to the following iterations:

- swap the root [largest element] with the last leaf
- remove the last leaf from the heap.
- heapify the root on the reduced heap.

The average and worst-case performance are

```
def heapsort(A_in):
A = np.copy(A_in)
# build a max heap
size = build_max_heap(A)
for i in range(size - 1, 0, -1):
# swap the root (largest element) with the last leaf
A[i], A[0] = A[0], A[i]
# removing largest element from the heap
size -= 1
# call _max_heapify from the root on the heap with remaining elements
max_heapify(A, size, node_idx=0)
return A
```

```
n = 10
A = rng.random(n, dtype=np.float64)
A
```

```
array([0.18525118, 0.99313433, 0.78561885, 0.44814329, 0.85044505,
0.86088208, 0.96716993, 0.17096352, 0.69956773, 0.8288503 ])
```

```
A_sorted_python = heapsort(A)
A_sorted_python
```

```
array([0.17096352, 0.18525118, 0.44814329, 0.69956773, 0.78561885,
0.8288503 , 0.85044505, 0.86088208, 0.96716993, 0.99313433])
```

```
A_ref = np.sort(A)
np.testing.assert_array_equal(A_sorted_python, A_ref)
```

Let's use Numba to make this Python code running fast.

## Numba

```
@njit
def max_heapify_numba(A, size, node_idx=0):
largest = node_idx
left_child = 2 * node_idx + 1
right_child = 2 * (node_idx + 1)
if (left_child < size) and (A[left_child] > A[largest]):
largest = left_child
if (right_child < size) and (A[right_child] > A[largest]):
largest = right_child
if largest != node_idx:
A[node_idx], A[largest] = A[largest], A[node_idx] # exchange 2 nodes
max_heapify_numba(A, size, largest)
@njit
def heapsort_numba(A_in):
A = np.copy(A_in)
# build a max heap
size = len(A)
node_idx = size // 2 - 1 # last non-leaf node index
for i in range(node_idx, -1, -1):
max_heapify_numba(A, size, node_idx=i)
for i in range(size - 1, 0, -1):
# swap the root (largest element) with the last leaf
A[i], A[0] = A[0], A[i]
# removing largest element from the heap
size -= 1
# call _max_heapify from the root on the heap with remaining elements
max_heapify_numba(A, size)
return A
```

```
A_sorted_numba = heapsort_numba(A)
A_sorted_numba
```

```
array([0.17096352, 0.18525118, 0.44814329, 0.69956773, 0.78561885,
0.8288503 , 0.85044505, 0.86088208, 0.96716993, 0.99313433])
```

`np.testing.assert_array_equal(A_sorted_numba, A_ref)`

## Cython

```
%%cython --compile-args=-Ofast
# cython: boundscheck=False, initializedcheck=False, wraparound=False
import cython
import numpy as np
from cython import ssize_t, double
@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _max_heapify(A: double[::1], size: ssize_t, node_idx: ssize_t) -> cython.void:
largest: ssize_t = node_idx
left_child = 2 * node_idx + 1
right_child = 2 * (node_idx + 1)
if left_child < size and A[left_child] > A[largest]:
largest = left_child
if right_child < size and A[right_child] > A[largest]:
largest = right_child
if largest != node_idx:
A[node_idx], A[largest] = A[largest], A[node_idx]
_max_heapify(A, size, largest)
@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _heapsort(A: double[::1]) -> cython.void:
i: cython.int
size: ssize_t = len(A)
node_idx: ssize_t = size // 2 - 1
for i in range(node_idx, -1, -1):
_max_heapify(A, size, i)
for i in range(size - 1, 0, -1):
A[i], A[0] = A[0], A[i]
size -= 1
_max_heapify(A, size, 0)
@cython.ccall
def heapsort_cython(A_in):
A = np.copy(A_in)
_heapsort(A)
return A
```

```
A_sorted_cython = heapsort_cython(A)
A_sorted_cython
```

```
array([0.17096352, 0.18525118, 0.44814329, 0.69956773, 0.78561885,
0.8288503 , 0.85044505, 0.86088208, 0.96716993, 0.99313433])
```

`np.testing.assert_array_equal(A_sorted_cython, A_ref)`

## Performance comparison

We do not include the Python version and compare the Numba and Cython versions with the NumPy *heapsort* implementation. It seems that this NumPy *heapsort* is written in C++ and is fairly well optimized.

```
out = perfplot.bench(
setup=lambda n: rng.random(n, dtype=np.float64),
kernels=[
lambda A: heapsort_numba(A),
lambda A: heapsort_cython(A),
lambda A: np.sort(A, kind="heapsort"),
],
labels=["Numba", "Cython", "NumPy"],
n_range=[10**k for k in range(1, 9)],
)
```

`out`

┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ n ┃ Numba ┃ Cython ┃ NumPy ┃ ┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩ │ 10 │ 1.266e-06 │ 1.819e-06 │ 1.7500000000000002e-06 │ │ 100 │ 5.587e-06 │ 5.004000000000001e-06 │ 3.0430000000000003e-06 │ │ 1000 │ 0.000116266 │ 6.1561e-05 │ 6.2518e-05 │ │ 10000 │ 0.001819482 │ 0.0008792880000000001 │ 0.000939073 │ │ 100000 │ 0.028507956 │ 0.014659915 │ 0.013501216000000002 │ │ 1000000 │ 0.42401333100000005 │ 0.26564063200000004 │ 0.224838018 │ │ 10000000 │ 5.938197069 │ 5.285706523 │ 3.4547337810000003 │ │ 100000000 │ 78.544810064 │ 84.04674386900001 │ 49.002946636000004 │ └───────────┴─────────────────────┴───────────────────────┴────────────────────────┘

```
figsize = (14, 14)
labels = out.labels
ms = 10
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1)
plt.loglog(
out.n_range,
out.n_range * np.log(out.n_range) * 5.0e-8,
"o-",
label="$c \; n \; ln(n)$",
)
for i, label in enumerate(labels):
plt.loglog(out.n_range, out.timings_s[i], "o-", ms=ms, label=label)
markers = cycle(("o", "v", "^", "<", ">", "s", "p", "P", "*", "h", "X", "D", "."))
for i, line in enumerate(ax.get_lines()):
marker = next(markers)
line.set_marker(marker)
plt.legend()
plt.grid("on")
_ = ax.set(
title="Timing comparison between Numba, Cython and NumPy heapsort",
xlabel="Array length (log scale)",
ylabel="Elapsed_time [s] (log scale)",
)
```

## Conclusion

Going from Python to Numba is seamless and is allowing us to reach a similar level of efficiency as with Cython [that turned 20 years old last week, happy birthday!!]. We can observe that the NumPy *heapsort* implementation is faster on larger arrays, but we do not know which optimizations did they implement.

## References

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2009. Introduction to Algorithms, Third Edition [3rd. ed.]. The MIT Press.

[2] Robert Sedgewick. 2002. Algorithms in C [3rd. ed.]. Addison-Wesley Longman Publishing Co., Inc., USA.