1
2
3
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 any) {
27 *h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
28 return
29 }
30
31 func (h *myHeap) Push(v any) {
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)
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)
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)
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 }
215
View as plain text