[Go] How to deal with sort and heap (2)

1. Quick Recap

In the previous post, we explored the concept of interfaces and how the sort package utilizes them. Understanding the abstract role that interfaces play, independent of specific implementations, was crucial. This abstraction allowed us to work with various types by tying them to the same interface, enabling them to be passed as parameters seamlessly.

By leveraging interfaces, we were able to use utility functions with multiple types. In the case of sorting, this was accomplished by implementing the required methods of the sort.Interface: Less, Len, and Swap. These methods enabled custom types to be sorted using the built-in sort functionality, providing flexibility and reusability in our code.

Now, building on that understanding, we’ll turn our attention to the heap package, which similarly uses interfaces to provide a powerful mechanism for managing heaps and priority queues. Like sorting, heap operations require us to implement a set of methods that define the behavior of the heap. These include Len, Less, Swap, as well as additional methods like Push and Pop that allow for dynamic modifications to the heap structure.

2. Heap package

Just like the sort package, Go’s heap package is built on the foundation of interfaces. Specifically, the heap.Interface provides the blueprint for creating and manipulating heaps. Heaps are commonly used for priority queues, where elements are organized based on their priority, typically using a binary heap structure.

To work with the heap package, you need to implement the following methods from the heap.Interface:

type Interface interface {
    sort.Interface            // Embeds sort.Interface
    Push(x interface{})       
    Pop() interface{}         
}

The heap.Interface extends sort.Interface by embedding it, meaning that in addition to the Len(), Less(), and Swap() methods that we used in the sort package, you must also implement two additional methods: Push() and Pop().

  • Len() int: Returns the number of elements in the heap.
  • Less(i, j int) bool: Defines the heap order by comparing elements at index i and j. For a min-heap, it returns true if the element at i is smaller than the element at j. For a max-heap, the opposite is true.
  • Swap(i, j int): Swaps the elements at indices i and j.
  • Push(x interface{}): Adds an element to the heap while maintaining heap properties.
  • Pop() interface{}: Removes and returns the root element (the highest or lowest priority element) from the heap.

2-a. Implementing heap interface

Let’s say we want to manage a heap of int values. We would start by implementing the heap.Interface for a custom type. Below is an example where we implement a min-heap of integers:

package main

import (
    "container/heap"
    "fmt"
)

// IntHeap is a min-heap of ints
type IntHeap []int

// Implement sort.Interface methods for IntHeap
func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

// Implement Push and Pop for heap.Interface
func (h *IntHeap) Push(x interface{}) {
    // Push uses a pointer receiver because we're modifying the slice in-place
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func main() {
    h := &IntHeap{2, 1, 5}
    heap.Init(h)        // Initialize the heap

    heap.Push(h, 3)     // Push an element
    fmt.Printf("Min: %d\n", (*h)[0]) // Min is the root element

    fmt.Printf("Removed: %d\n", heap.Pop(h)) // Pop removes the root element
    fmt.Printf("Heap: %v\n", *h)             // Print the current state of the heap
}

Implementing sort interface is not a big deal, however, there is a difference with Push and Pop methods. To implement Push and Pop, you have to append the element at the end of the slice. Why is it?

source: https://medium.com/@verdi/binary-heap-basics-40a0f3f41c8f

It is because heap is maintained by bubbling the element up and down. When we push the element, it is added at the last node of the heap and it is bubbled up. By doing this, we can maintain the heap structure without rotating process like AVL tree.

source: https://medium.com/@verdi/binary-heap-basics-40a0f3f41c8f
source: https://medium.com/@verdi/binary-heap-basics-40a0f3f41c8f

When we pop an element, the root element is removed and the last element is moved to the root. It is then bubbled down to restore the heap property.

Beside these basic logic of push and pop, there is one more thing to keep our eyes on. It is that receivers of push and pop are pointer types. Why is this?

2-b. Pointer receiver

When we implement a method, we can use either a value receiver or a pointer receiver. With a value receiver, the instance is copied, so any changes made within the method do not affect the original instance.

You might wonder about the Swap method in this context. It seems like Swap changes elements within the slice, right? That’s a great observation. However, the key point is that the slice itself is a reference type—its underlying array is stored elsewhere in memory, and the slice holds a pointer to it. So even though we are using a value receiver, the elements in the slice can still be modified because the slice contains a reference to the array, not the actual data itself.

However, Push and Pop must use pointer receivers because they affect the slice’s underlying fields. While appending an element to a slice might seem to work fine with a value receiver, changes to the len and cap fields of the slice must also be reflected in the original heap. Using a value receiver would only update a copy, leaving the original slice unchanged. If you want to know more about Go slice, take a look at this documentation.

3. Using the heap package

After implementing the heap.Interface for our custom type, such as IntHeap, the next step is to properly initialize and manage the heap using the heap package. This involves a few critical steps that ensure the heap maintains its structure and works efficiently for operations like inserting and removing elements.

3-a. Heap functions

To heapify a heap interface, we have to initialize it by calling Init function of heap package.

h := &IntHeap{2, 1, 5}
heap.Init(h)  

Default setting of int heap is min heap. If you want make an max heap, you have to fix Less method.

Let’s look at how Push and Pop work in a Max-Heap. The Push method adds a new element to the heap and ensures that the heap property is maintained by “bubbling up” the element to its correct position. The Pop method removes the largest element (the root) from the heap, replaces it with the last element, and then “bubbles down” that element to restore the heap property.

package main

import (
    "container/heap"
    "fmt"
)

// IntHeap is a Max-Heap of ints
type IntHeap []int

// Implement sort.Interface methods for IntHeap
func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }  // Max-Heap comparison
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

// Implement Push and Pop for heap.Interface
func (h *IntHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func main() {
    // Initialize a max-heap
    h := &IntHeap{2, 1, 5}
    heap.Init(h)

    // Push an element into the heap
    heap.Push(h, 3)
    fmt.Printf("Max after Push: %d\n", (*h)[0]) // Should print the largest element in the heap

    // Pop the largest element from the heap
    maxElem := heap.Pop(h)
    fmt.Printf("Popped: %d\n", maxElem)
    fmt.Printf("Heap after Pop: %v\n", *h) // Should print the remaining elements in the heap
}

Push Example:

  • After calling heap.Push(h, 3), the value 3 is added to the heap. Since this is a Max-Heap, the heap will rearrange itself so that the largest value is at the root. For example, if 5 is still the largest value after adding 3, it will remain at the root.

Pop Example:

  • When heap.Pop(h) is called, the largest element (the root of the Max-Heap) is removed. The last element of the heap is moved to the root, and then it “bubbles down” to its correct position to maintain the heap property.

4. Conclusion

In this post, we explored the heap package in Go and how to effectively implement and use heaps. We started by understanding how to implement the heap.Interface to create custom heaps, including both Min-Heap and Max-Heap variations.

We then delved into the practical implementation of Init, Push and Pop methods, seeing how elements can be dynamically added and removed while maintaining the heap structure. The use of pointer receivers was emphasized, particularly for methods that modify the underlying slice, ensuring that changes to the heap persist.

By grasping these concepts, we now have a solid understanding of how heaps work in Go, and how they can be applied in various real-world scenarios such as priority queues, algorithm optimization, and efficient data management. This knowledge equips us to tackle more complex problems using Go’s powerful heap functionalities.

Leave a Reply

Your email address will not be published. Required fields are marked *