/
IntervalTree.go
98 lines (85 loc) · 2.11 KB
/
IntervalTree.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
package nsqd
import (
"github.com/Workiva/go-datastructures/augmentedtree"
)
type QueueInterval struct {
start int64
end int64
endCnt int64
}
func (QI *QueueInterval) Start() int64 {
return QI.start
}
func (QI *QueueInterval) End() int64 {
return QI.end
}
func (QI *QueueInterval) EndCnt() int64 {
return QI.endCnt
}
// the augmentedtree use the low and the id to determin if the interval is the duplicate
// so here we use the end as the id of segment
func (QI *QueueInterval) ID() uint64 {
return uint64(QI.end)
}
func (QI *QueueInterval) LowAtDimension(dim uint64) int64 {
return QI.start
}
func (QI *QueueInterval) HighAtDimension(dim uint64) int64 {
return QI.end
}
func (QI *QueueInterval) OverlapsAtDimension(inter augmentedtree.Interval, dim uint64) bool {
if inter.HighAtDimension(dim) < QI.start {
return false
}
if inter.LowAtDimension(dim) > QI.end {
return false
}
return true
}
type IntervalTree struct {
tr augmentedtree.Tree
}
func NewIntervalTree() *IntervalTree {
return &IntervalTree{
tr: augmentedtree.New(1),
}
}
func (IT *IntervalTree) AddOrMerge(inter *QueueInterval) *QueueInterval {
overlaps := IT.tr.Query(inter)
if len(overlaps) == 1 && overlaps[0].LowAtDimension(0) <= inter.LowAtDimension(0) &&
overlaps[0].HighAtDimension(0) >= inter.HighAtDimension(0) {
return overlaps[0].(*QueueInterval)
} else if 0 == len(overlaps) {
IT.tr.Add(inter)
return inter
} else {
qi := &QueueInterval{}
qi.start = inter.Start()
qi.end = inter.End()
qi.endCnt = inter.EndCnt()
for _, v := range overlaps {
if v.LowAtDimension(0) < qi.start {
qi.start = v.LowAtDimension(0)
}
if v.LowAtDimension(0) > qi.end {
qi.end = v.HighAtDimension(0)
qi.endCnt = v.(*QueueInterval).EndCnt()
}
IT.tr.Delete(v)
}
return qi
}
return nil
}
func (IT *IntervalTree) Len() int64 {
return int64(IT.tr.Len())
}
func (self *IntervalTree) DeleteInterval(inter *QueueInterval) {
overlaps := self.tr.Query(inter)
for _, v := range overlaps {
if v.LowAtDimension(1) == inter.LowAtDimension(1) &&
v.HighAtDimension(1) == inter.HighAtDimension(1) {
self.tr.Delete(v)
}
}
}