Home » Blogs » tung's blog

Max-heap in Go

21Nov2009
tung's picture

I was helping somebody on a forum a few days ago with a C++ max-heap implementation. I thought it might be fun to try it out in Go.

package main
 
import (
    "container/vector";
    "flag";
    "rand";
    "time";
)
 
var num = flag.Int("n", 10, "number of elements to heap sort")
 
func main() {
    flag.Parse();
    rand.Seed(time.Nanoseconds());
    v := vector.NewIntVector(*num);
    for i := 1; i <= *num; i++ {
        v.Set(i - 1, rand.Int() % 100)
    }
    heapSort(v);
    printHeap(v, 0, 0);
}
 
func left(i int) int {
    return i * 2 + 1
}
 
func right(i int) int {
    return i * 2 + 2
}
 
func parent(i int) int {
    return (i - 1) / 2
}
 
// Swap vec[i] with its parent if its larger. Stop at the top.
func maxHeapify(vec *vector.IntVector, i int) {
    if i <= 0 {
        return
    }
    p := parent(i);
    if vec.At(p) < vec.At(i) {
        temp := vec.At(p);
        vec.Set(p, vec.At(i));
        vec.Set(i, temp);
        maxHeapify(vec, p);
    }
}
 
// Given a list of numbers in any order, max heap sort it.
func heapSort(vec *vector.IntVector) {
    for i := 0; i < vec.Len(); i++ {
        maxHeapify(vec, i)
    }
}
 
func printHeap(vec *vector.IntVector, i int, nest int) {
    if i >= vec.Len() {
        return
    }
    for n := 0; n < nest; n++ {
        print("    ")
    }
    print(vec.At(i));
    print("\n");
    printHeap(vec, left(i), nest + 1);
    printHeap(vec, right(i), nest + 1);
}

The things to pay attention to are the left, right and parent functions. Most descriptions in textbooks and on the web use 1-based arrays, so they have simpler formulas. Since arrays in Go are zero-based, the formulas had to be tweaked so the children and parents lined up correctly.

You'll also notice that there's no pop or push. The original person's exercise was just the sorting of an array (it's a vector here) into a heap-like arrangement. If I wanted to push, it'd look something like:

func push(heap *vector.IntVector, x int) {
    heap.Push(x);
    maxHeapify(heap, heap.Len() - 1);
}

A pop is a bit more involved. I would pluck the top element out, take the left/right replacement for it, take the left/right replacement for that, etc. until I hit the bottom, where I would pull in the last vector element and do the usual heap sort algorithm at that point. That last step needs to be done, since there's no guarantee that the leaf in its new position maintains the heap ordering property.

func pop(heap *vector.IntVector) (x int) {
    x = heap[0];
    current := 0;
    l, r := left(current), right(current);
    for l < heap.Len() && r < heap.Len() {
        if (heap[l] > heap[r]) {
            heap[current] = heap[l];
            current = l;
        } else {
            heap[current] = heap[r];
            current = r;
        }
        l, r = left(current), right(current);
    }
    // Handle case where there's just a left child.
    if l < heap.Len() {
        heap[current] = heap[l];
        current = l;
    }
    // heap.Pop() below is a vector pop, not a true heap pop, so it removes the end element.
    heap[current] = heap.Pop();
    maxHeapify(heap, current);
}

Because of the way heaps are packed into arrays/vectors, it's never possible for any node to have a right child, but no left child.

I learned this 3 years ago in a second year data structures course, haven't used it once, and I still pulled this out of the top of my head. A relief for sure, but I'm afraid that if I don't do any serious coding soon, it'll all vanish into the ether.