diff options
Diffstat (limited to 'libgo/go/container/heap')
-rw-r--r-- | libgo/go/container/heap/heap.go | 102 | ||||
-rw-r--r-- | libgo/go/container/heap/heap_test.go | 166 |
2 files changed, 268 insertions, 0 deletions
diff --git a/libgo/go/container/heap/heap.go b/libgo/go/container/heap/heap.go new file mode 100644 index 000000000..4435a57c4 --- /dev/null +++ b/libgo/go/container/heap/heap.go @@ -0,0 +1,102 @@ +// Copyright 2009 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// This package provides heap operations for any type that implements +// heap.Interface. +// +package heap + +import "sort" + +// Any type that implements heap.Interface may be used as a +// min-heap with the following invariants (established after +// Init has been called): +// +// !h.Less(j, i) for 0 <= i < h.Len() and j = 2*i+1 or 2*i+2 and j < h.Len() +// +type Interface interface { + sort.Interface + Push(x interface{}) + Pop() interface{} +} + + +// A heaper must be initialized before any of the heap operations +// can be used. Init is idempotent with respect to the heap invariants +// and may be called whenever the heap invariants may have been invalidated. +// Its complexity is O(n) where n = h.Len(). +// +func Init(h Interface) { + // heapify + n := h.Len() + for i := n/2 - 1; i >= 0; i-- { + down(h, i, n) + } +} + + +// Push pushes the element x onto the heap. The complexity is +// O(log(n)) where n = h.Len(). +// +func Push(h Interface, x interface{}) { + h.Push(x) + up(h, h.Len()-1) +} + + +// Pop removes the minimum element (according to Less) from the heap +// and returns it. The complexity is O(log(n)) where n = h.Len(). +// Same as Remove(h, 0). +// +func Pop(h Interface) interface{} { + n := h.Len() - 1 + h.Swap(0, n) + down(h, 0, n) + return h.Pop() +} + + +// Remove removes the element at index i from the heap. +// The complexity is O(log(n)) where n = h.Len(). +// +func Remove(h Interface, i int) interface{} { + n := h.Len() - 1 + if n != i { + h.Swap(i, n) + down(h, i, n) + up(h, i) + } + return h.Pop() +} + + +func up(h Interface, j int) { + for { + i := (j - 1) / 2 // parent + if i == j || h.Less(i, j) { + break + } + h.Swap(i, j) + j = i + } +} + + +func down(h Interface, i, n int) { + for { + j1 := 2*i + 1 + if j1 >= n { + 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(i, j) { + break + } + h.Swap(i, j) + i = j + } +} diff --git a/libgo/go/container/heap/heap_test.go b/libgo/go/container/heap/heap_test.go new file mode 100644 index 000000000..89d444dd5 --- /dev/null +++ b/libgo/go/container/heap/heap_test.go @@ -0,0 +1,166 @@ +// Copyright 2009 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package heap + +import ( + "testing" + "container/vector" +) + + +type myHeap struct { + // A vector.Vector implements sort.Interface except for Less, + // and it implements Push and Pop as required for heap.Interface. + vector.Vector +} + + +func (h *myHeap) Less(i, j int) bool { return h.At(i).(int) < h.At(j).(int) } + + +func (h *myHeap) verify(t *testing.T, i int) { + n := h.Len() + j1 := 2*i + 1 + j2 := 2*i + 2 + if j1 < n { + if h.Less(j1, i) { + t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h.At(i), j1, h.At(j1)) + return + } + h.verify(t, j1) + } + if j2 < n { + if h.Less(j2, i) { + t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h.At(i), j1, h.At(j2)) + return + } + h.verify(t, j2) + } +} + + +func TestInit0(t *testing.T) { + h := new(myHeap) + for i := 20; i > 0; i-- { + h.Push(0) // all elements are the same + } + Init(h) + h.verify(t, 0) + + for i := 1; h.Len() > 0; i++ { + x := Pop(h).(int) + h.verify(t, 0) + if x != 0 { + t.Errorf("%d.th pop got %d; want %d", i, x, 0) + } + } +} + + +func TestInit1(t *testing.T) { + h := new(myHeap) + for i := 20; i > 0; i-- { + h.Push(i) // all elements are different + } + Init(h) + h.verify(t, 0) + + for i := 1; h.Len() > 0; i++ { + x := Pop(h).(int) + h.verify(t, 0) + if x != i { + t.Errorf("%d.th pop got %d; want %d", i, x, i) + } + } +} + + +func Test(t *testing.T) { + h := new(myHeap) + h.verify(t, 0) + + for i := 20; i > 10; i-- { + h.Push(i) + } + Init(h) + h.verify(t, 0) + + for i := 10; i > 0; i-- { + Push(h, i) + h.verify(t, 0) + } + + for i := 1; h.Len() > 0; i++ { + x := Pop(h).(int) + if i < 20 { + Push(h, 20+i) + } + h.verify(t, 0) + if x != i { + t.Errorf("%d.th pop got %d; want %d", i, x, i) + } + } +} + + +func TestRemove0(t *testing.T) { + h := new(myHeap) + for i := 0; i < 10; i++ { + h.Push(i) + } + h.verify(t, 0) + + for h.Len() > 0 { + i := h.Len() - 1 + x := Remove(h, i).(int) + if x != i { + t.Errorf("Remove(%d) got %d; want %d", i, x, i) + } + h.verify(t, 0) + } +} + + +func TestRemove1(t *testing.T) { + h := new(myHeap) + for i := 0; i < 10; i++ { + h.Push(i) + } + h.verify(t, 0) + + for i := 0; h.Len() > 0; i++ { + x := Remove(h, 0).(int) + if x != i { + t.Errorf("Remove(0) got %d; want %d", x, i) + } + h.verify(t, 0) + } +} + + +func TestRemove2(t *testing.T) { + N := 10 + + h := new(myHeap) + for i := 0; i < N; i++ { + h.Push(i) + } + h.verify(t, 0) + + m := make(map[int]bool) + for h.Len() > 0 { + m[Remove(h, (h.Len()-1)/2).(int)] = true + h.verify(t, 0) + } + + if len(m) != N { + t.Errorf("len(m) = %d; want %d", len(m), N) + } + for i := 0; i < len(m); i++ { + if !m[i] { + t.Errorf("m[%d] doesn't exist", i) + } + } +} |