forked from nsqio/nsq
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pqueue.go
75 lines (63 loc) · 1.26 KB
/
pqueue.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
package pqueue
import (
"container/heap"
)
type Item struct {
Value interface{}
Priority int64
Index int
}
// this is a priority queue as implemented by a min heap
// ie. the 0th element is the *lowest* value
type PriorityQueue []*Item
func New(capacity int) PriorityQueue {
return make(PriorityQueue, 0, capacity)
}
func (pq PriorityQueue) Len() int {
return len(pq)
}
func (pq PriorityQueue) Less(i, j int) bool {
return pq[i].Priority < pq[j].Priority
}
func (pq PriorityQueue) Swap(i, j int) {
pq[i], pq[j] = pq[j], pq[i]
pq[i].Index = i
pq[j].Index = j
}
func (pq *PriorityQueue) Push(x interface{}) {
n := len(*pq)
c := cap(*pq)
if n+1 > c {
npq := make(PriorityQueue, n, c*2)
copy(npq, *pq)
*pq = npq
}
*pq = (*pq)[0 : n+1]
item := x.(*Item)
item.Index = n
(*pq)[n] = item
}
func (pq *PriorityQueue) Pop() interface{} {
n := len(*pq)
c := cap(*pq)
if n < (c/2) && c > 25 {
npq := make(PriorityQueue, n, c/2)
copy(npq, *pq)
*pq = npq
}
item := (*pq)[n-1]
item.Index = -1
*pq = (*pq)[0 : n-1]
return item
}
func (pq *PriorityQueue) PeekAndShift(max int64) (*Item, int64) {
if pq.Len() == 0 {
return nil, 0
}
item := (*pq)[0]
if item.Priority > max {
return nil, item.Priority - max
}
heap.Remove(pq, 0)
return item, 0
}