forked from rodrigo-brito/ninjabot
-
Notifications
You must be signed in to change notification settings - Fork 0
/
priorityqueue.go
117 lines (102 loc) · 1.8 KB
/
priorityqueue.go
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package model
import "sync"
type PriorityQueue struct {
sync.Mutex
length int
data []Item
notifyCallbacks []func(Item)
}
type Item interface {
Less(Item) bool
}
func NewPriorityQueue(data []Item) *PriorityQueue {
q := &PriorityQueue{}
q.data = data
q.length = len(data)
if q.length > 0 {
i := q.length >> 1
for ; i >= 0; i-- {
q.down(i)
}
}
return q
}
func (q *PriorityQueue) Push(item Item) {
q.Lock()
defer q.Unlock()
q.data = append(q.data, item)
q.length++
q.up(q.length - 1)
for _, notify := range q.notifyCallbacks {
go notify(item)
}
}
func (q *PriorityQueue) PopLock() <-chan Item {
ch := make(chan Item)
q.notifyCallbacks = append(q.notifyCallbacks, func(_ Item) {
ch <- q.Pop()
})
return ch
}
func (q *PriorityQueue) Pop() Item {
q.Lock()
defer q.Unlock()
if q.length == 0 {
return nil
}
top := q.data[0]
q.length--
if q.length > 0 {
q.data[0] = q.data[q.length]
q.down(0)
}
q.data = q.data[:len(q.data)-1]
return top
}
func (q *PriorityQueue) Peek() Item {
q.Lock()
defer q.Unlock()
if q.length == 0 {
return nil
}
return q.data[0]
}
func (q *PriorityQueue) Len() int {
q.Lock()
defer q.Unlock()
return q.length
}
func (q *PriorityQueue) down(pos int) {
data := q.data
halfLength := q.length >> 1
item := data[pos]
for pos < halfLength {
left := (pos << 1) + 1
right := left + 1
best := data[left]
if right < q.length && data[right].Less(best) {
left = right
best = data[right]
}
if !best.Less(item) {
break
}
data[pos] = best
pos = left
}
data[pos] = item
}
func (q *PriorityQueue) up(pos int) {
data := q.data
item := data[pos]
for pos > 0 {
parent := (pos - 1) >> 1
current := data[parent]
if !item.Less(current) {
break
}
data[pos] = current
pos = parent
}
data[pos] = item
}