-
Notifications
You must be signed in to change notification settings - Fork 81
/
priority_queue.go
80 lines (66 loc) · 1.42 KB
/
priority_queue.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
package sync
import (
"container/heap"
)
type (
key = string
)
type item struct {
key string
priority int64
index int
}
type priorityQueue struct {
items []*item
itemByKey map[string]*item
}
func (pq *priorityQueue) isPending(key key) bool {
if _, ok := pq.itemByKey[key]; ok {
return true
}
return false
}
func (pq *priorityQueue) add(key key, priority int64) {
if _, ok := pq.itemByKey[key]; ok {
return
}
heap.Push(pq, &item{key: key, priority: priority})
}
func (pq *priorityQueue) remove(key key) {
if item, ok := pq.itemByKey[key]; ok {
heap.Remove(pq, item.index)
delete(pq.itemByKey, key)
}
}
func (pq *priorityQueue) pop() *item {
item, _ := heap.Pop(pq).(*item)
return item
}
func (pq *priorityQueue) peek() *item {
return pq.items[0]
}
func (pq priorityQueue) Len() int { return len(pq.items) }
func (pq priorityQueue) Less(i, j int) bool {
return pq.items[i].priority < pq.items[j].priority
}
func (pq priorityQueue) Swap(i, j int) {
pq.items[i], pq.items[j] = pq.items[j], pq.items[i]
pq.items[i].index = i
pq.items[j].index = j
}
func (pq *priorityQueue) Push(x interface{}) {
n := len(pq.items)
item, _ := x.(*item)
item.index = n
pq.items = append(pq.items, item)
pq.itemByKey[item.key] = item
}
func (pq *priorityQueue) Pop() interface{} {
old := pq.items
n := len(old)
item := old[n-1]
item.index = -1
pq.items = old[0 : n-1]
delete(pq.itemByKey, item.key)
return item
}