summaryrefslogtreecommitdiff
path: root/pkg/priority_queue/binary_heap.go
blob: 00c7386978accd0eb2957656cdf4e738db73414d (plain)
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
/*
binary heap (min-heap) algorithm used as a core for the priority queue
*/

package priorityqueue

type BinHeap []PQItem

func NewBinHeap() *BinHeap {
	return &BinHeap{}
}

func (bh *BinHeap) Init(items []PQItem) {
	arraySize := len(items) - 1

	for i := arraySize/2 - 1; i >= 0; i-- {
		bh.fixDown(i, arraySize)
	}

	for i := arraySize - 1; i >= 1; i-- {
		items[0], items[i] = items[i], items[0]
		bh.fixDown(0, i-1)
	}
}

func (bh *BinHeap) fixUp() {
	k := len(*bh) - 1
	p := (k - 1) >> 1 // k-1 / 2

	for k > 0 {
		cur, par := (*bh)[k], (*bh)[p]

		if cur.Priority() < par.Priority() {
			bh.swap(k, p)
			k = p
			p = (k - 1) >> 1
		} else {
			return
		}
	}
}

func (bh *BinHeap) swap(i, j int) {
	(*bh)[i], (*bh)[j] = (*bh)[j], (*bh)[i]
}

func (bh *BinHeap) fixDown(curr, end int) {
	cOneIdx := curr*2 + 1
	for cOneIdx <= end {
		cTwoIdx := -1
		if curr*2+2 <= end {
			cTwoIdx = curr*2 + 2
		}

		idxToSwap := cOneIdx
		if cTwoIdx > -1 && (*bh)[cTwoIdx].Priority() < (*bh)[cOneIdx].Priority() {
			idxToSwap = cTwoIdx
		}
		if (*bh)[idxToSwap].Priority() < (*bh)[curr].Priority() {
			bh.swap(curr, idxToSwap)
			curr = idxToSwap
			cOneIdx = curr*2 + 1
		} else {
			return
		}
	}
}

func (bh *BinHeap) Insert(item PQItem) {
	*bh = append(*bh, item)
	bh.fixUp()
}

func (bh *BinHeap) GetMax() PQItem {
	l := len(*bh)
	if l == 0 {
		return nil
	}

	bh.swap(0, l-1)

	item := (*bh)[l-1]
	*bh = (*bh)[0 : l-1]
	bh.fixDown(0, l-2)
	return item
}