DSGO/heap/binary/heap.go at master · PeterRK/DSGO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
package binary

import (
	"cmp"
)

//二叉堆，底层采用数组。
//Build的复杂度为O(N)，Top的复杂度为O(1)，其余核心操作复杂度为O(logN)。
type Heap[T cmp.Ordered] struct {
	vec []T
}

func (hp *Heap[T]) Size() int {
	return len(hp.vec)
}

func (hp *Heap[T]) IsEmpty() bool {
	return len(hp.vec) == 0
}

func (hp *Heap[T]) Clear() {
	clear(hp.vec)
	hp.vec = hp.vec[:0]
}

func (hp *Heap[T]) Top() T {
	if hp.IsEmpty() {
		panic("empty heap")
	}
	return hp.vec[0]
}

func (hp *Heap[T]) BuildInPlace(list []T) {
	size := len(list)
	hp.vec = list
	for idx := size/2 - 1; idx >= 0; idx-- {
		hp.down(idx)
	}
}

func (hp *Heap[T]) Build(list []T) {
	vec := make([]T, len(list))
	copy(vec, list)
	hp.BuildInPlace(vec)
}

func (hp *Heap[T]) Push(unit T) {
	place := len(hp.vec)
	hp.vec = append(hp.vec, unit)
	hp.up(place)
}

func (hp *Heap[T]) Pop() T {
	size := hp.Size()
	if size == 0 {
		panic("empty heap")
	}
	unit := hp.vec[0]
	if size == 1 {
		var zero T
		hp.vec[0] = zero
		hp.vec = hp.vec[:0]
	} else {
		hp.vec[0] = hp.vec[size-1]
		var zero T
		hp.vec[size-1] = zero
		hp.vec = hp.vec[:size-1]
		hp.down(0)
	}
	return unit
}

func (hp *Heap[T]) down(pos int) {
	tgt := hp.vec[pos]
	kid, last := pos*2+1, len(hp.vec)-1
	for kid < last {
		if hp.vec[kid+1] < hp.vec[kid] {
			kid++
		}
		if tgt <= hp.vec[kid] {
			break
		}
		hp.vec[pos] = hp.vec[kid]
		pos, kid = kid, kid*2+1
	}
	if kid == last && tgt > hp.vec[kid] {
		hp.vec[pos], pos = hp.vec[kid], kid
	}
	hp.vec[pos] = tgt
}

func (hp *Heap[T]) up(pos int) {
	tgt := hp.vec[pos]
	for pos > 0 {
		parent := (pos - 1) / 2
		if hp.vec[parent] <= tgt {
			break
		}
		hp.vec[pos], pos = hp.vec[parent], parent
	}
	hp.vec[pos] = tgt
}