Skip to content

Commit a55e1fd

Browse files
FlorimondFlorimond
committed
Adds generic binary heap and d-ary heap
1 parent 567f6a3 commit a55e1fd

File tree

4 files changed

+356
-0
lines changed

4 files changed

+356
-0
lines changed

‎heap/binary.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
MIT License
3+
4+
Copyright (c) 2021 Florimond Husquinet
5+
6+
Permission is hereby granted, free of charge, to any person obtaining a copy
7+
of this software and associated documentation files (the "Software"), to deal
8+
in the Software without restriction, including without limitation the rights
9+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+
copies of the Software, and to permit persons to whom the Software is
11+
furnished to do so, subject to the following conditions:
12+
13+
The above copyright notice and this permission notice shall be included in all
14+
copies or substantial portions of the Software.
15+
16+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+
SOFTWARE.
23+
*/
24+
25+
/* A generic implementation of a binary heap */
26+
package heap
27+
28+
type Heap[T any] struct {
29+
data []T
30+
compare func(T, T) int
31+
}
32+
33+
func NewHeap[T any](compare func(T, T) int) *Heap[T] {
34+
return &Heap[T]{
35+
data: make([]T, 0),
36+
compare: compare,
37+
}
38+
}
39+
40+
func (h *Heap[T]) Len() int {
41+
return len(h.data)
42+
}
43+
44+
func (h *Heap[T]) Push(value T) {
45+
h.data = append(h.data, value)
46+
h.bubbleUp(h.Len() - 1)
47+
}
48+
49+
func (h *Heap[T]) Pop() (ok bool, value T) {
50+
if h.Len() == 0 {
51+
return false, value
52+
}
53+
var top = h.data[0]
54+
h.data[0] = h.data[h.Len()-1]
55+
h.data = h.data[:h.Len()-1]
56+
h.sinkDown(0)
57+
return true, top
58+
}
59+
60+
// Min heap: if a node is less than its parent, swap them.
61+
func (h *Heap[T]) bubbleUp(index int) {
62+
if index == 0 {
63+
return
64+
}
65+
var parent = (index - 1) / 2
66+
if h.compare(h.data[index], h.data[parent]) < 0 {
67+
h.swap(index, parent)
68+
h.bubbleUp(parent)
69+
}
70+
}
71+
72+
// Min heap: if a node is greater than its children, swap the node with the smallest child.
73+
func (h *Heap[T]) sinkDown(index int) {
74+
var left = index*2 + 1
75+
var right = index*2 + 2
76+
var smallest = index
77+
if left < h.Len() && h.compare(h.data[left], h.data[smallest]) < 0 {
78+
smallest = left
79+
}
80+
if right < h.Len() && h.compare(h.data[right], h.data[smallest]) < 0 {
81+
smallest = right
82+
}
83+
if smallest != index {
84+
h.swap(index, smallest)
85+
h.sinkDown(smallest)
86+
}
87+
}
88+
89+
func (h *Heap[T]) swap(i, j int) {
90+
var tmp = h.data[i]
91+
h.data[i] = h.data[j]
92+
h.data[j] = tmp
93+
}

‎heap/binary_test.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package heap
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/assert"
7+
)
8+
9+
func TestHeap(t *testing.T) {
10+
h := NewHeap(func(a, b int) int {
11+
if a < b {
12+
return -1
13+
}
14+
if a > b {
15+
return 1
16+
}
17+
return 0
18+
})
19+
20+
h.Push(10)
21+
h.Push(15)
22+
h.Push(1)
23+
h.Push(5)
24+
h.Push(9)
25+
h.Push(7)
26+
h.Push(2)
27+
28+
ok, value := h.Pop()
29+
assert.True(t, ok)
30+
assert.Equal(t, 1, value)
31+
32+
ok, value = h.Pop()
33+
assert.True(t, ok)
34+
assert.Equal(t, 2, value)
35+
36+
ok, value = h.Pop()
37+
assert.True(t, ok)
38+
assert.Equal(t, 5, value)
39+
40+
ok, value = h.Pop()
41+
assert.True(t, ok)
42+
assert.Equal(t, 7, value)
43+
44+
ok, value = h.Pop()
45+
assert.True(t, ok)
46+
assert.Equal(t, 9, value)
47+
48+
ok, value = h.Pop()
49+
assert.True(t, ok)
50+
assert.Equal(t, 10, value)
51+
52+
ok, value = h.Pop()
53+
assert.True(t, ok)
54+
assert.Equal(t, 15, value)
55+
56+
ok, _ = h.Pop()
57+
assert.False(t, ok)
58+
}
59+
60+
func BenchmarkBinaryHeap(b *testing.B) {
61+
h := NewHeap(func(a, b int) int {
62+
if a < b {
63+
return -1
64+
}
65+
if a > b {
66+
return 1
67+
}
68+
return 0
69+
})
70+
for i := 0; i < b.N; i++ {
71+
h.Push(b.N)
72+
}
73+
for i := 0; i < b.N; i++ {
74+
h.Pop()
75+
}
76+
}

‎heap/dary.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
/*
2+
MIT License
3+
4+
Copyright (c) 2021 Florimond Husquinet
5+
6+
Permission is hereby granted, free of charge, to any person obtaining a copy
7+
of this software and associated documentation files (the "Software"), to deal
8+
in the Software without restriction, including without limitation the rights
9+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+
copies of the Software, and to permit persons to whom the Software is
11+
furnished to do so, subject to the following conditions:
12+
13+
The above copyright notice and this permission notice shall be included in all
14+
copies or substantial portions of the Software.
15+
16+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+
SOFTWARE.
23+
*/
24+
25+
/*
26+
A generic implementation of a d-ary heap.
27+
28+
The d-ary heap or d-heap is a priority queue data structure, a generalization
29+
of the binary heap in which the nodes have d children instead of 2.
30+
*/
31+
package heap
32+
33+
import "math"
34+
35+
type DaryHeap[T any] struct {
36+
d int
37+
data []T
38+
compare func(T, T) int
39+
}
40+
41+
func NewDaryHeap[T any](d int, compare func(T, T) int) *DaryHeap[T] {
42+
return &DaryHeap[T]{
43+
d: d,
44+
data: make([]T, 0),
45+
compare: compare,
46+
}
47+
}
48+
49+
func (h *DaryHeap[T]) Len() int {
50+
return len(h.data)
51+
}
52+
53+
func (h *DaryHeap[T]) Push(value T) {
54+
h.data = append(h.data, value)
55+
h.bubbleUp(h.Len() - 1)
56+
}
57+
58+
func (h *DaryHeap[T]) Pop() (ok bool, value T) {
59+
if h.Len() == 0 {
60+
return false, value
61+
}
62+
var top = h.data[0]
63+
h.data[0] = h.data[h.Len()-1]
64+
h.data = h.data[:h.Len()-1]
65+
h.sinkDown(0)
66+
return true, top
67+
}
68+
69+
// Min heap: if a node is less than its parent, swap them.
70+
func (h *DaryHeap[T]) bubbleUp(index int) {
71+
if index == 0 {
72+
return
73+
}
74+
var parent = (index - 1) / h.d // Todo: make test fail if d is not 2 but you divide by 2
75+
if h.compare(h.data[index], h.data[parent]) < 0 {
76+
h.swap(index, parent)
77+
h.bubbleUp(parent)
78+
}
79+
}
80+
81+
// Min heap: if a node is greater than its children, swap the node with the smallest child.
82+
func (h *DaryHeap[T]) sinkDown(index int) {
83+
var childrenIndex = int(math.Pow(2, float64(index)))
84+
var smallest = index
85+
86+
for i := 0; i < h.d; i++ {
87+
var child = childrenIndex + i
88+
if child >= h.Len() {
89+
break
90+
}
91+
if h.compare(h.data[child], h.data[smallest]) < 0 {
92+
smallest = child
93+
}
94+
}
95+
if smallest != index {
96+
h.swap(index, smallest)
97+
h.sinkDown(smallest)
98+
}
99+
}
100+
101+
func (h *DaryHeap[T]) swap(i, j int) {
102+
var tmp = h.data[i]
103+
h.data[i] = h.data[j]
104+
h.data[j] = tmp
105+
}

‎heap/dary_test.go‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package heap
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/assert"
7+
)
8+
9+
func TestDaryHeap(t *testing.T) {
10+
h := NewDaryHeap(6, func(a, b int) int {
11+
if a < b {
12+
return -1
13+
}
14+
if a > b {
15+
return 1
16+
}
17+
return 0
18+
})
19+
20+
h.Push(10)
21+
h.Push(1)
22+
h.Push(5)
23+
h.Push(9)
24+
h.Push(7)
25+
h.Push(2)
26+
h.Push(15)
27+
28+
h.Push(20)
29+
30+
ok, value := h.Pop()
31+
assert.True(t, ok)
32+
assert.Equal(t, 1, value)
33+
34+
ok, value = h.Pop()
35+
assert.True(t, ok)
36+
assert.Equal(t, 2, value)
37+
38+
ok, value = h.Pop()
39+
assert.True(t, ok)
40+
assert.Equal(t, 5, value)
41+
42+
ok, value = h.Pop()
43+
assert.True(t, ok)
44+
assert.Equal(t, 7, value)
45+
46+
ok, value = h.Pop()
47+
assert.True(t, ok)
48+
assert.Equal(t, 9, value)
49+
50+
ok, value = h.Pop()
51+
assert.True(t, ok)
52+
assert.Equal(t, 10, value)
53+
54+
ok, value = h.Pop()
55+
assert.True(t, ok)
56+
assert.Equal(t, 15, value)
57+
58+
ok, value = h.Pop()
59+
assert.True(t, ok)
60+
assert.Equal(t, 20, value)
61+
62+
ok, _ = h.Pop()
63+
assert.False(t, ok)
64+
}
65+
66+
func BenchmarkDaryHeap(b *testing.B) {
67+
h := NewDaryHeap(3, func(a, b int) int {
68+
if a < b {
69+
return -1
70+
}
71+
if a > b {
72+
return 1
73+
}
74+
return 0
75+
})
76+
for i := 0; i < b.N; i++ {
77+
h.Push(b.N)
78+
}
79+
for i := 0; i < b.N; i++ {
80+
h.Pop()
81+
}
82+
}

0 commit comments

Comments
 (0)