extensible Heap in Go

heap.go

package heap                                            

// HeapData has to be implemented by any data
// that will be inserted into the Heap Data structure.
// You can create MinHeap or MaxHeap.
type HeapData interface {                           
        IsLessThan(h HeapData) bool                 
}                                                   

//Heap holds the entire heap object
type Heap struct {                 
        elems     []HeapData       
        size      int              
        isMinHeap bool             
}                                  

//NewMinHeap creates a Heap with Minimum Value at the root
func NewMinHeap() Heap {                                  
        return Heap{make([]HeapData, 0, 100), 0, true}    
}                                                         

//NewMaxHeap creates a Heap with Maximum Value at the root
func NewMaxHeap() Heap {                                  
        return Heap{make([]HeapData, 0, 100), 0, false}   
}                                                         

//parent is i, children are in 2i+1 and 2i+2
//child is i, parent is i/2 (integer division)
func (h *Heap) heapifyDown() {                
        if h.size > 0 {                       
                curNode := 0                  
                for curNode < h.size {        
                        leftChild := 2*curNode + 1
                        rightChild := 2*curNode + 2
                        if leftChild >= h.size {   
                                break              
                        }                          
                        selectChild := leftChild   
                        if rightChild < h.size && h.isMinHeap && h.elems[rightChild].IsLessThan(h.elems[leftChild]) {
                                selectChild = rightChild                                                             
                        }                                                                                            
                        if rightChild < h.size && !h.isMinHeap && h.elems[leftChild].IsLessThan(h.elems[rightChild]) {
                                selectChild = rightChild                                                              
                        }                                                                                             
                        if (h.isMinHeap && h.elems[selectChild].IsLessThan(h.elems[curNode])) ||                      
                                (!h.isMinHeap && h.elems[curNode].IsLessThan(h.elems[selectChild])) {                 
                                t := h.elems[curNode]                                                                 
                                h.elems[curNode] = h.elems[selectChild]                                               
                                h.elems[selectChild] = t                                                              
                                curNode = selectChild                                                                 
                        } else {                                                                                      
                                break                                                                                 
                        }                                                                                             
                }                                                                                                     
        }                                                                                                             
}                                                                                                                     

func (h *Heap) heapifyUp() {
        if h.size > 0 {     
                curNode := h.size - 1
                for curNode != 0 {   
                        parent := curNode / 2
                        if (h.isMinHeap && h.elems[curNode].IsLessThan(h.elems[parent])) ||
                                (!h.isMinHeap && h.elems[parent].IsLessThan(h.elems[curNode])) {
                                t := h.elems[curNode]                                           
                                h.elems[curNode] = h.elems[parent]                              
                                h.elems[parent] = t                                             
                                curNode = parent                                                
                        } else {                                                                
                                break                                                           
                        }                                                                       
                }                                                                               
        }                                                                                       
}                                                                                               

//Insert inserts an object into the heap object
func (h *Heap) Insert(val HeapData) {          
        h.elems = append(h.elems, val)
        h.size++
        h.heapifyUp()
}

//Peek returns the root of the heap object
//is either the smallest or largest - depending on MinHeap or MaxHeap
//If heap is empty will return ok=false
func (h *Heap) Peek() (val HeapData, ok bool) {
        if h.size > 0 {
                return h.elems[0], true
        }
        return nil, false
}

//ExtractTop Removes and returns the root of the heap object
//is either the smallest or largest - depending on MinHeap or MaxHeap
//If heap is empty will return ok=false
func (h *Heap) ExtractTop() (val HeapData, ok bool) {
        if h.size > 0 {
                ret := h.elems[0]
                h.elems[0] = h.elems[h.size-1]
                h.elems = h.elems[:h.size-1]
                h.size--
                h.heapifyDown()
                return ret, true
        }
        return nil, false
}