Skip to content

Instantly share code, notes, and snippets.

@helinwang
Created July 19, 2018 04:57
Show Gist options
  • Save helinwang/c9dbd0a6938cd6755b9ff636ce54e790 to your computer and use it in GitHub Desktop.
Save helinwang/c9dbd0a6938cd6755b9ff636ce54e790 to your computer and use it in GitHub Desktop.
a heap implementation in Go
package types
// Heap is a max heap
type Heap struct {
buf []int
}
func (h *Heap) Peak() int {
return h.buf[0]
}
func (h *Heap) up(idx int) {
if idx == 0 {
return
}
parent := (idx - 1) / 2
if h.buf[parent] < h.buf[idx] {
h.buf[parent], h.buf[idx] = h.buf[idx], h.buf[parent]
h.up(parent)
return
}
}
func (h *Heap) Insert(key int) {
h.buf = append(h.buf, key)
h.up(len(h.buf) - 1)
}
func (h *Heap) down(idx int) {
leftChild := idx*2 + 1
rightChild := idx*2 + 2
if leftChild >= len(h.buf) {
return
}
swap := -1
if h.buf[idx] < h.buf[leftChild] {
if rightChild >= len(h.buf) || h.buf[leftChild] >= h.buf[rightChild] {
swap = leftChild
} else {
swap = rightChild
}
} else if rightChild < len(h.buf) && h.buf[idx] < h.buf[rightChild] {
swap = rightChild
}
if swap < 0 {
return
}
h.buf[swap], h.buf[idx] = h.buf[idx], h.buf[swap]
h.down(swap)
}
func (h *Heap) Delete() int {
max := h.buf[0]
h.buf[0], h.buf[len(h.buf)-1] = h.buf[len(h.buf)-1], h.buf[0]
h.buf = h.buf[:len(h.buf)-1]
h.down(0)
return max
}
func (h *Heap) Build() {
for i := len(h.buf) - 1; i >= 0; i -= 2 {
h.down((i - 1) / 2)
}
}
// put the following to heap_test.go
package types
import (
"math/rand"
"testing"
"github.com/stretchr/testify/assert"
)
func TestHeap(t *testing.T) {
const total = 101
h := &Heap{}
perm := rand.Perm(total)
for _, v := range perm {
h.Insert(v)
}
assert.Equal(t, total-1, h.Peak())
for i := total - 1; i >= 0; i-- {
assert.Equal(t, i, h.Delete())
}
}
func TestBuildHeap(t *testing.T) {
const total = 101
perm := rand.Perm(total)
h := &Heap{buf: perm}
h.Build()
assert.Equal(t, total-1, h.Peak())
for i := total - 1; i >= 0; i-- {
assert.Equal(t, i, h.Delete())
}
}
func TestHeapSort(t *testing.T) {
const total = 101
perm := rand.Perm(total)
h := &Heap{buf: perm}
h.Build()
assert.Equal(t, total-1, h.Peak())
for i := total - 1; i >= 0; i-- {
assert.Equal(t, i, h.Delete())
}
for i := 0; i < total; i++ {
assert.Equal(t, i, perm[i])
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment