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 n := h.Len() 37 j1 := 2*i + 1 38 j2 := 2*i + 2 39 if j1 < n { 40 if h.Less(j1, i) { 41 t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j1]) 42 return 43 } 44 h.verify(t, j1) 45 } 46 if j2 < n { 47 if h.Less(j2, i) { 48 t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j2]) 49 return 50 } 51 h.verify(t, j2) 52 } 53 } 54 55 func TestInit0(t *testing.T) { 56 h := new(myHeap) 57 for i := 20; i > 0; i-- { 58 h.Push(0) // all elements are the same 59 } 60 Init(h) 61 h.verify(t, 0) 62 63 for i := 1; h.Len() > 0; i++ { 64 x := Pop(h).(int) 65 h.verify(t, 0) 66 if x != 0 { 67 t.Errorf("%d.th pop got %d; want %d", i, x, 0) 68 } 69 } 70 } 71 72 func TestInit1(t *testing.T) { 73 h := new(myHeap) 74 for i := 20; i > 0; i-- { 75 h.Push(i) // all elements are different 76 } 77 Init(h) 78 h.verify(t, 0) 79 80 for i := 1; h.Len() > 0; i++ { 81 x := Pop(h).(int) 82 h.verify(t, 0) 83 if x != i { 84 t.Errorf("%d.th pop got %d; want %d", i, x, i) 85 } 86 } 87 } 88 89 func Test(t *testing.T) { 90 h := new(myHeap) 91 h.verify(t, 0) 92 93 for i := 20; i > 10; i-- { 94 h.Push(i) 95 } 96 Init(h) 97 h.verify(t, 0) 98 99 for i := 10; i > 0; i-- { 100 Push(h, i) 101 h.verify(t, 0) 102 } 103 104 for i := 1; h.Len() > 0; i++ { 105 x := Pop(h).(int) 106 if i < 20 { 107 Push(h, 20+i) 108 } 109 h.verify(t, 0) 110 if x != i { 111 t.Errorf("%d.th pop got %d; want %d", i, x, i) 112 } 113 } 114 } 115 116 func TestRemove0(t *testing.T) { 117 h := new(myHeap) 118 for i := 0; i < 10; i++ { 119 h.Push(i) 120 } 121 h.verify(t, 0) 122 123 for h.Len() > 0 { 124 i := h.Len() - 1 125 x := Remove(h, i).(int) 126 if x != i { 127 t.Errorf("Remove(%d) got %d; want %d", i, x, i) 128 } 129 h.verify(t, 0) 130 } 131 } 132 133 func TestRemove1(t *testing.T) { 134 h := new(myHeap) 135 for i := 0; i < 10; i++ { 136 h.Push(i) 137 } 138 h.verify(t, 0) 139 140 for i := 0; h.Len() > 0; i++ { 141 x := Remove(h, 0).(int) 142 if x != i { 143 t.Errorf("Remove(0) got %d; want %d", x, i) 144 } 145 h.verify(t, 0) 146 } 147 } 148 149 func TestRemove2(t *testing.T) { 150 N := 10 151 152 h := new(myHeap) 153 for i := 0; i < N; i++ { 154 h.Push(i) 155 } 156 h.verify(t, 0) 157 158 m := make(map[int]bool) 159 for h.Len() > 0 { 160 m[Remove(h, (h.Len()-1)/2).(int)] = true 161 h.verify(t, 0) 162 } 163 164 if len(m) != N { 165 t.Errorf("len(m) = %d; want %d", len(m), N) 166 } 167 for i := 0; i < len(m); i++ { 168 if !m[i] { 169 t.Errorf("m[%d] doesn't exist", i) 170 } 171 } 172 } 173 174 func BenchmarkDup(b *testing.B) { 175 const n = 10000 176 h := make(myHeap, n) 177 for i := 0; i < b.N; i++ { 178 for j := 0; j < n; j++ { 179 Push(&h, 0) // all elements are the same 180 } 181 for h.Len() > 0 { 182 Pop(&h) 183 } 184 } 185 } 186 187 func TestFix(t *testing.T) { 188 h := new(myHeap) 189 h.verify(t, 0) 190 191 for i := 200; i > 0; i -= 10 { 192 Push(h, i) 193 } 194 h.verify(t, 0) 195 196 if (*h)[0] != 10 { 197 t.Fatalf("Expected head to be 10, was %d", (*h)[0]) 198 } 199 (*h)[0] = 210 200 Fix(h, 0) 201 h.verify(t, 0) 202 203 for i := 100; i > 0; i-- { 204 elem := rand.Intn(h.Len()) 205 if i&1 == 0 { 206 (*h)[elem] *= 2 207 } else { 208 (*h)[elem] /= 2 209 } 210 Fix(h, elem) 211 h.verify(t, 0) 212 } 213 } 214
