libgo/go/container/heap/heap_test.go

   1 // Copyright 2009 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package heap
   6
   7 import (
   8         "math/rand"
   9         "testing"
  10 )
  11
  12 type myHeap []int
  13
  14 func (h *myHeap) Less(i, j int) bool {
  15         return (*h)[i] < (*h)[j]
  16 }
  17
  18 func (h *myHeap) Swap(i, j int) {
  19         (*h)[i], (*h)[j] = (*h)[j], (*h)[i]
  20 }
  21
  22 func (h *myHeap) Len() int {
  23         return len(*h)
  24 }
  25
  26 func (h *myHeap) Pop() (v interface{}) {
  27         *h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
  28         return
  29 }
  30
  31 func (h *myHeap) Push(v interface{}) {
  32         *h = append(*h, v.(int))
  33 }
  34
  35 func (h myHeap) verify(t *testing.T, i int) {
  36         t.Helper()
  37         n := h.Len()
  38         j1 := 2*i + 1
  39         j2 := 2*i + 2
  40         if j1 < n {
  41                 if h.Less(j1, i) {
  42                         t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j1])
  43                         return
  44                 }
  45                 h.verify(t, j1)
  46         }
  47         if j2 < n {
  48                 if h.Less(j2, i) {
  49                         t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j2])
  50                         return
  51                 }
  52                 h.verify(t, j2)
  53         }
  54 }
  55
  56 func TestInit0(t *testing.T) {
  57         h := new(myHeap)
  58         for i := 20; i > 0; i-- {
  59                 h.Push(0) // all elements are the same
  60         }
  61         Init(h)
  62         h.verify(t, 0)
  63
  64         for i := 1; h.Len() > 0; i++ {
  65                 x := Pop(h).(int)
  66                 h.verify(t, 0)
  67                 if x != 0 {
  68                         t.Errorf("%d.th pop got %d; want %d", i, x, 0)
  69                 }
  70         }
  71 }
  72
  73 func TestInit1(t *testing.T) {
  74         h := new(myHeap)
  75         for i := 20; i > 0; i-- {
  76                 h.Push(i) // all elements are different
  77         }
  78         Init(h)
  79         h.verify(t, 0)
  80
  81         for i := 1; h.Len() > 0; i++ {
  82                 x := Pop(h).(int)
  83                 h.verify(t, 0)
  84                 if x != i {
  85                         t.Errorf("%d.th pop got %d; want %d", i, x, i)
  86                 }
  87         }
  88 }
  89
  90 func Test(t *testing.T) {
  91         h := new(myHeap)
  92         h.verify(t, 0)
  93
  94         for i := 20; i > 10; i-- {
  95                 h.Push(i)
  96         }
  97         Init(h)
  98         h.verify(t, 0)
  99
 100         for i := 10; i > 0; i-- {
 101                 Push(h, i)
 102                 h.verify(t, 0)
 103         }
 104
 105         for i := 1; h.Len() > 0; i++ {
 106                 x := Pop(h).(int)
 107                 if i < 20 {
 108                         Push(h, 20+i)
 109                 }
 110                 h.verify(t, 0)
 111                 if x != i {
 112                         t.Errorf("%d.th pop got %d; want %d", i, x, i)
 113                 }
 114         }
 115 }
 116
 117 func TestRemove0(t *testing.T) {
 118         h := new(myHeap)
 119         for i := 0; i < 10; i++ {
 120                 h.Push(i)
 121         }
 122         h.verify(t, 0)
 123
 124         for h.Len() > 0 {
 125                 i := h.Len() - 1
 126                 x := Remove(h, i).(int)
 127                 if x != i {
 128                         t.Errorf("Remove(%d) got %d; want %d", i, x, i)
 129                 }
 130                 h.verify(t, 0)
 131         }
 132 }
 133
 134 func TestRemove1(t *testing.T) {
 135         h := new(myHeap)
 136         for i := 0; i < 10; i++ {
 137                 h.Push(i)
 138         }
 139         h.verify(t, 0)
 140
 141         for i := 0; h.Len() > 0; i++ {
 142                 x := Remove(h, 0).(int)
 143                 if x != i {
 144                         t.Errorf("Remove(0) got %d; want %d", x, i)
 145                 }
 146                 h.verify(t, 0)
 147         }
 148 }
 149
 150 func TestRemove2(t *testing.T) {
 151         N := 10
 152
 153         h := new(myHeap)
 154         for i := 0; i < N; i++ {
 155                 h.Push(i)
 156         }
 157         h.verify(t, 0)
 158
 159         m := make(map[int]bool)
 160         for h.Len() > 0 {
 161                 m[Remove(h, (h.Len()-1)/2).(int)] = true
 162                 h.verify(t, 0)
 163         }
 164
 165         if len(m) != N {
 166                 t.Errorf("len(m) = %d; want %d", len(m), N)
 167         }
 168         for i := 0; i < len(m); i++ {
 169                 if !m[i] {
 170                         t.Errorf("m[%d] doesn't exist", i)
 171                 }
 172         }
 173 }
 174
 175 func BenchmarkDup(b *testing.B) {
 176         const n = 10000
 177         h := make(myHeap, 0, n)
 178         for i := 0; i < b.N; i++ {
 179                 for j := 0; j < n; j++ {
 180                         Push(&h, 0) // all elements are the same
 181                 }
 182                 for h.Len() > 0 {
 183                         Pop(&h)
 184                 }
 185         }
 186 }
 187
 188 func TestFix(t *testing.T) {
 189         h := new(myHeap)
 190         h.verify(t, 0)
 191
 192         for i := 200; i > 0; i -= 10 {
 193                 Push(h, i)
 194         }
 195         h.verify(t, 0)
 196
 197         if (*h)[0] != 10 {
 198                 t.Fatalf("Expected head to be 10, was %d", (*h)[0])
 199         }
 200         (*h)[0] = 210
 201         Fix(h, 0)
 202         h.verify(t, 0)
 203
 204         for i := 100; i > 0; i-- {
 205                 elem := rand.Intn(h.Len())
 206                 if i&1 == 0 {
 207                         (*h)[elem] *= 2
 208                 } else {
 209                         (*h)[elem] /= 2
 210                 }
 211                 Fix(h, elem)
 212                 h.verify(t, 0)
 213         }
 214 }