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
5package heap
6
7import (
8	"math/rand"
9	"testing"
10)
11
12type myHeap []int
13
14func (h *myHeap) Less(i, j int) bool {
15	return (*h)[i] < (*h)[j]
16}
17
18func (h *myHeap) Swap(i, j int) {
19	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
20}
21
22func (h *myHeap) Len() int {
23	return len(*h)
24}
25
26func (h *myHeap) Pop() (v any) {
27	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
28	return
29}
30
31func (h *myHeap) Push(v any) {
32	*h = append(*h, v.(int))
33}
34
35func (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
56func 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
73func 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
90func 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
117func 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
134func 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
150func 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
175func 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
188func 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