Skip to content

Instantly share code, notes, and snippets.

@rogeriomarques

rogeriomarques/gheap.go

Forked from Komosa/gheap.go
Created November 1, 2015 13:38
Show Gist options
  • Save rogeriomarques/66b626b51d294c3c51ee to your computer and use it in GitHub Desktop.
Save rogeriomarques/66b626b51d294c3c51ee to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. @Komosa Komosa revised this gist Nov 1, 2015. No changes.
  2. @Komosa Komosa created this gist Nov 1, 2015.
    61 changes: 61 additions & 0 deletions gheap.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,61 @@
    // this package is mostly copy-pasted from golang's std container/heap
    // I changed Interface, Pop and Push in order to get rid of type assertions
    // usage can be found in test file, benchmarks show that we can get about 27% of speedup for 1000 elems, or 10% for 10M elems (tested on go1.4)
    package gheap

    type Interface interface {
    Len() int
    Less(i, j int) bool
    Swap(i, j int)
    }

    func Init(h Interface) {
    // heapify
    n := h.Len()
    for i := n/2 - 1; i >= 0; i-- {
    down(h, i, n)
    }
    }

    // put new element at Len() position before calling Push(h)
    // you must also increase Len
    func Push(h Interface) {
    up(h, h.Len()-1)
    }

    // after calling Pop() top element will be at position Len()-1
    // you must decrease Len after this call
    func Pop(h Interface) {
    n := h.Len() - 1
    h.Swap(0, n)
    down(h, 0, n)
    }

    func up(h Interface, j int) {
    for {
    i := (j - 1) / 2 // parent
    if i == j || !h.Less(j, i) {
    break
    }
    h.Swap(i, j)
    j = i
    }
    }

    func down(h Interface, i, n int) {
    for {
    j1 := 2*i + 1
    if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
    break
    }
    j := j1 // left child
    if j2 := j1 + 1; j2 < n && !h.Less(j1, j2) {
    j = j2 // = 2*i + 2 // right child
    }
    if !h.Less(j, i) {
    break
    }
    h.Swap(i, j)
    i = j
    }
    }
    151 changes: 151 additions & 0 deletions gheap_test.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,151 @@
    package gheap

    import (
    "container/heap"
    "fmt"
    "math/rand"
    "testing"
    )

    const N = 10000000

    var ints []int
    var out1, out2 []int

    func init() {
    for i := 0; i < N; i++ {
    ints = append(ints, rand.Int())
    }
    out1 = make([]int, N)
    out2 = make([]int, N)
    }

    func UseTypeAssertion(n int) {
    mh := myHeap(make([]int, 0, n))
    h := &mh

    a, b := n/2, n/4
    c := a + b

    for i := 0; i < a; i++ {
    heap.Push(h, ints[i])
    }
    for i := 0; i < b; i++ {
    out1[i] = heap.Pop(h).(int)
    }
    for i := a; i < c; i++ {
    heap.Push(h, ints[i])
    }
    for i := b; i < a; i++ {
    out1[i] = heap.Pop(h).(int)
    }
    for i := c; i < n; i++ {
    heap.Push(h, ints[i])
    }
    for i := a; i < n; i++ {
    out1[i] = heap.Pop(h).(int)
    }
    }

    func NotUseTypeAssertion(n int) {
    mh := myGHeap(make([]int, 0, n))
    h := &mh

    a, b := n/2, n/4
    c := a + b

    for i := 0; i < a; i++ {
    h.Push(ints[i])
    }
    for i := 0; i < b; i++ {
    out2[i] = h.Pop()
    }
    for i := a; i < c; i++ {
    h.Push(ints[i])
    }
    for i := b; i < a; i++ {
    out2[i] = h.Pop()
    }
    for i := c; i < n; i++ {
    h.Push(ints[i])
    }
    for i := a; i < n; i++ {
    out2[i] = h.Pop()
    }
    }

    func areEq(n int) bool {
    for i := 0; i < n; i++ {
    if out1[i] != out2[i] {
    fmt.Println("mismatch!")
    return false
    }
    }
    return true
    }

    func benchmark(f func(int), n int, b *testing.B) {
    for i := 0; i < b.N; i++ {
    f(n)
    }
    }

    func BenchmarkStdHeap1000(b *testing.B) { benchmark(UseTypeAssertion, 1000, b) }
    func BenchmarkMyGHeap1000(b *testing.B) { benchmark(NotUseTypeAssertion, 1000, b) }
    func BenchmarkCheckEq1000(b *testing.B) { areEq(1000) }
    func BenchmarkStdHeap10000(b *testing.B) { benchmark(UseTypeAssertion, 10000, b) }
    func BenchmarkMyGHeap10000(b *testing.B) { benchmark(NotUseTypeAssertion, 10000, b) }
    func BenchmarkCheckEq10000(b *testing.B) { areEq(10000) }
    func BenchmarkStdHeap100000(b *testing.B) { benchmark(UseTypeAssertion, 100000, b) }
    func BenchmarkMyGHeap100000(b *testing.B) { benchmark(NotUseTypeAssertion, 100000, b) }
    func BenchmarkCheckEq100000(b *testing.B) { areEq(100000) }
    func BenchmarkStdHeap1000000(b *testing.B) { benchmark(UseTypeAssertion, 1000000, b) }
    func BenchmarkMyGHeap1000000(b *testing.B) { benchmark(NotUseTypeAssertion, 1000000, b) }
    func BenchmarkCheckEq1000000(b *testing.B) { areEq(1000000) }
    func BenchmarkStdHeap10000000(b *testing.B) { benchmark(UseTypeAssertion, 10000000, b) }
    func BenchmarkMyGHeap10000000(b *testing.B) { benchmark(NotUseTypeAssertion, 10000000, b) }
    func BenchmarkCheckEq10000000(b *testing.B) { areEq(10000000) }

    // for without type assertion version
    type myGHeap []int

    func (h *myGHeap) Less(i, j int) bool {
    return (*h)[i] < (*h)[j]
    }
    func (h *myGHeap) Swap(i, j int) {
    (*h)[i], (*h)[j] = (*h)[j], (*h)[i]
    }
    func (h *myGHeap) Len() int {
    return len(*h)
    }
    func (h *myGHeap) Pop() (v int) {
    //~ gheap.Pop(h)
    Pop(h)
    *h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
    return
    }
    func (h *myGHeap) Push(v int) {
    *h = append(*h, v)
    //~ gheap.Push(h)
    Push(h)
    }

    // for type assertion version
    type myHeap []int

    func (h *myHeap) Less(i, j int) bool {
    return (*h)[i] < (*h)[j]
    }
    func (h *myHeap) Swap(i, j int) {
    (*h)[i], (*h)[j] = (*h)[j], (*h)[i]
    }
    func (h *myHeap) Len() int {
    return len(*h)
    }
    func (h *myHeap) Pop() (v interface{}) {
    *h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
    return
    }
    func (h *myHeap) Push(v interface{}) {
    *h = append(*h, v.(int))
    }