👋 Welcome Gophers! In this lesson, we are going to implement the heap sort algorithm that we looked at in the previous lesson.
Implementation
Let’s start by creating a new file called main.go in a directory called 04-heap-sort/. We’ll build up the heap sort implementation step by step.
Let’s start with the basic setup:
package main
import "fmt"
func main() {
fmt.Println("Heap Sort Algorithm in Go")
}
Now, let’s define the function signatures we’ll need:
// HeapSort - sorts the input slice using the heap sort algorithm
func HeapSort(input []int) {
}
// Heapify - maintains the max-heap property for a subtree rooted at index i
func Heapify(input []int, n int, i int) {
}
Let’s start with the Heapify function. This is the core operation that maintains the max-heap property:
func Heapify(input []int, n int, i int) {
largest := i
left := 2*i + 1 // left child
right := 2*i + 2 // right child
// if left child is larger than root
if left < n && input[left] > input[largest] {
largest = left
}
// if right child is larger than largest so far
if right < n && input[right] > input[largest] {
largest = right
}
// if largest is not root, swap and heapify recursively
if largest != i {
input[i], input[largest] = input[largest], input[i]
Heapify(input, n, largest)
}
}
The Heapify function works by:
- Finding the largest element among the node and its children
- If the largest is not the current node, swap them
- Recursively heapify the affected subtree
Now let’s implement the main HeapSort function:
func HeapSort(input []int) {
n := len(input)
// build max heap
// we start from the last non-leaf node which is at index (n/2 - 1)
for i := n/2 - 1; i >= 0; i-- {
Heapify(input, n, i)
}
// extract elements from heap one by one
for i := n - 1; i > 0; i-- {
// move current root (maximum) to the end
input[0], input[i] = input[i], input[0]
// heapify the reduced heap
Heapify(input, i, 0)
}
}
The HeapSort function has two phases:
- Build phase: Convert the array into a max-heap by calling Heapify for all non-leaf nodes
- Extraction phase: Repeatedly move the maximum element (at root) to the end and heapify the remaining heap
Now let’s test it out by creating a slice in our main function and calling our new HeapSort function:
func main() {
fmt.Println("Heap Sort Algorithm in Go")
unsortedInput := []int{5, 3, 4, 1, 2}
fmt.Println("Before:", unsortedInput)
HeapSort(unsortedInput)
fmt.Println("After:", unsortedInput)
}
Perfect! Let’s run this from the command line:
$ go run main.go
Heap Sort Algorithm in Go
Before: [5 3 4 1 2]
After: [1 2 3 4 5]
Excellent! Our heap sort is working correctly. Let’s add a version with debug output to see what’s happening:
func HeapSortDebug(input []int) {
n := len(input)
fmt.Println("Building max heap...")
for i := n/2 - 1; i >= 0; i-- {
HeapifyDebug(input, n, i, 0)
}
fmt.Println("After building heap:", input)
fmt.Println("\nExtracting elements...")
for i := n - 1; i > 0; i-- {
fmt.Printf("Moving %d to final position\n", input[0])
input[0], input[i] = input[i], input[0]
HeapifyDebug(input, i, 0, 0)
fmt.Printf("Heap state: %v\n", input[:i])
}
}
func HeapifyDebug(input []int, n int, i int, depth int) {
indent := ""
for j := 0; j < depth; j++ {
indent += " "
}
largest := i
left := 2*i + 1
right := 2*i + 2
if left < n && input[left] > input[largest] {
largest = left
}
if right < n && input[right] > input[largest] {
largest = right
}
if largest != i {
fmt.Printf("%sSwapping %d and %d\n", indent, input[i], input[largest])
input[i], input[largest] = input[largest], input[i]
HeapifyDebug(input, n, largest, depth+1)
}
}
When we run this with debug output, we can see the heapification process:
$ go run main.go
Heap Sort Algorithm in Go
Building max heap...
Swapping 1 and 5
Swapping 3 and 4
After building heap: [5 3 4 1 2]
Extracting elements...
Moving 5 to final position
Swapping 2 and 3
Heap state: [4 3 2]
Moving 4 to final position
Swapping 2 and 3
Heap state: [3 2]
Moving 3 to final position
Heap state: [2]
After: [1 2 3 4 5]
Perfect! You can see how the algorithm builds the heap and then extracts elements one by one.
Testing
Let’s create a proper test to validate our implementation:
func TestHeapSort(t *testing.T) {
tests := []struct {
input []int
expected []int
}{
{[]int{5, 3, 4, 1, 2}, []int{1, 2, 3, 4, 5}},
{[]int{1}, []int{1}},
{[]int{}, []int{}},
{[]int{2, 1}, []int{1, 2}},
{[]int{1, 2, 3, 4, 5}, []int{1, 2, 3, 4, 5}},
{[]int{5, 4, 3, 2, 1}, []int{1, 2, 3, 4, 5}},
{[]int{9, 5, 6, 2, 3, 7, 1, 4, 8}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9}},
}
for _, tt := range tests {
input := make([]int, len(tt.input))
copy(input, tt.input)
HeapSort(input)
if !slicesEqual(input, tt.expected) {
t.Errorf("HeapSort(%v) = %v, want %v", tt.input, input, tt.expected)
}
}
}
func slicesEqual(a, b []int) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
Why Heap Sort is Great
Heap sort is one of the most elegant sorting algorithms because it:
- Guarantees O(n log n): No worst-case scenarios like quicksort
- Uses O(1) space: Sorts in place with minimal memory overhead
- Is in-place: No need for temporary arrays
- Works well in practice: Despite being theoretically good, it has decent cache behavior
The main downside is that it’s not stable (doesn’t preserve the order of equal elements) and is slightly more complex to understand than simpler algorithms.
Conclusion
Awesome! We now have a working implementation of the heap sort algorithm in Go! This algorithm combines the best of both worlds: the guaranteed O(n log n) performance of merge sort with the O(1) space efficiency of insertion sort. It’s a powerful tool to have in your algorithm toolkit.
Quiz Time
Try your hand at these questions below to validate what we’ve just covered in this lesson: