diff options
author | Luke Shumaker <lukeshu@lukeshu.com> | 2023-02-12 02:44:35 -0700 |
---|---|---|
committer | Luke Shumaker <lukeshu@lukeshu.com> | 2023-02-12 02:44:35 -0700 |
commit | 128e4d9aa876e14a1203ce98bfaa7ad399ad97c7 (patch) | |
tree | 039c3c549414c21b15d58c3d695ee87c3feb1402 /lib/containers/intervaltree.go | |
parent | 53d7fbb73869eb5defa1ca5c52b26abd346b13b9 (diff) | |
parent | 696a7d192e5eefa53230168a4b200ec0560c8a10 (diff) |
Merge branch 'lukeshu/containers'
Diffstat (limited to 'lib/containers/intervaltree.go')
-rw-r--r-- | lib/containers/intervaltree.go | 136 |
1 files changed, 70 insertions, 66 deletions
diff --git a/lib/containers/intervaltree.go b/lib/containers/intervaltree.go index 16bc916..7b96526 100644 --- a/lib/containers/intervaltree.go +++ b/lib/containers/intervaltree.go @@ -4,81 +4,80 @@ package containers -type intervalKey[K Ordered[K]] struct { +type interval[K Ordered[K]] struct { Min, Max K } -func (ival intervalKey[K]) ContainsFn(fn func(K) int) bool { - return fn(ival.Min) >= 0 && fn(ival.Max) <= 0 -} - -func (a intervalKey[K]) Cmp(b intervalKey[K]) int { - if d := a.Min.Cmp(b.Min); d != 0 { +// Compare implements Ordered. +func (a interval[K]) Compare(b interval[K]) int { + if d := a.Min.Compare(b.Min); d != 0 { return d } - return a.Max.Cmp(b.Max) + return a.Max.Compare(b.Max) +} + +// ContainsFn returns whether this interval contains the range matched +// by the given function. +func (ival interval[K]) ContainsFn(fn func(K) int) bool { + return fn(ival.Min) >= 0 && fn(ival.Max) <= 0 } type intervalValue[K Ordered[K], V any] struct { - Val V - SpanOfChildren intervalKey[K] + Val V + ValSpan interval[K] + ChildSpan interval[K] +} + +// Compare implements Ordered. +func (a intervalValue[K, V]) Compare(b intervalValue[K, V]) int { + return a.ValSpan.Compare(b.ValSpan) } type IntervalTree[K Ordered[K], V any] struct { MinFn func(V) K MaxFn func(V) K - inner RBTree[intervalKey[K], intervalValue[K, V]] -} - -func (t *IntervalTree[K, V]) keyFn(v intervalValue[K, V]) intervalKey[K] { - return intervalKey[K]{ - Min: t.MinFn(v.Val), - Max: t.MaxFn(v.Val), - } + inner RBTree[intervalValue[K, V]] } func (t *IntervalTree[K, V]) attrFn(node *RBNode[intervalValue[K, V]]) { - max := t.MaxFn(node.Value.Val) - if node.Left != nil && node.Left.Value.SpanOfChildren.Max.Cmp(max) > 0 { - max = node.Left.Value.SpanOfChildren.Max + max := node.Value.ValSpan.Max + if node.Left != nil && node.Left.Value.ChildSpan.Max.Compare(max) > 0 { + max = node.Left.Value.ChildSpan.Max } - if node.Right != nil && node.Right.Value.SpanOfChildren.Max.Cmp(max) > 0 { - max = node.Right.Value.SpanOfChildren.Max + if node.Right != nil && node.Right.Value.ChildSpan.Max.Compare(max) > 0 { + max = node.Right.Value.ChildSpan.Max } - node.Value.SpanOfChildren.Max = max + node.Value.ChildSpan.Max = max - min := t.MinFn(node.Value.Val) - if node.Left != nil && node.Left.Value.SpanOfChildren.Min.Cmp(min) < 0 { - min = node.Left.Value.SpanOfChildren.Min + min := node.Value.ValSpan.Min + if node.Left != nil && node.Left.Value.ChildSpan.Min.Compare(min) < 0 { + min = node.Left.Value.ChildSpan.Min } - if node.Right != nil && node.Right.Value.SpanOfChildren.Min.Cmp(min) < 0 { - min = node.Right.Value.SpanOfChildren.Min + if node.Right != nil && node.Right.Value.ChildSpan.Min.Compare(min) < 0 { + min = node.Right.Value.ChildSpan.Min } - node.Value.SpanOfChildren.Min = min + node.Value.ChildSpan.Min = min } func (t *IntervalTree[K, V]) init() { - if t.inner.KeyFn == nil { - t.inner.KeyFn = t.keyFn + if t.inner.AttrFn == nil { t.inner.AttrFn = t.attrFn } } -func (t *IntervalTree[K, V]) Delete(min, max K) { - t.init() - t.inner.Delete(intervalKey[K]{ - Min: min, - Max: max, - }) -} - func (t *IntervalTree[K, V]) Equal(u *IntervalTree[K, V]) bool { return t.inner.Equal(&u.inner) } func (t *IntervalTree[K, V]) Insert(val V) { t.init() - t.inner.Insert(intervalValue[K, V]{Val: val}) + t.inner.Insert(intervalValue[K, V]{ + Val: val, + ValSpan: interval[K]{ + Min: t.MinFn(val), + Max: t.MaxFn(val), + }, + }) } func (t *IntervalTree[K, V]) Min() (K, bool) { @@ -86,7 +85,7 @@ func (t *IntervalTree[K, V]) Min() (K, bool) { var zero K return zero, false } - return t.inner.root.Value.SpanOfChildren.Min, true + return t.inner.root.Value.ChildSpan.Min, true } func (t *IntervalTree[K, V]) Max() (K, bool) { @@ -94,22 +93,18 @@ func (t *IntervalTree[K, V]) Max() (K, bool) { var zero K return zero, false } - return t.inner.root.Value.SpanOfChildren.Max, true -} - -func (t *IntervalTree[K, V]) Lookup(k K) (V, bool) { - return t.Search(k.Cmp) + return t.inner.root.Value.ChildSpan.Max, true } func (t *IntervalTree[K, V]) Search(fn func(K) int) (V, bool) { node := t.inner.root for node != nil { switch { - case t.keyFn(node.Value).ContainsFn(fn): + case node.Value.ValSpan.ContainsFn(fn): return node.Value.Val, true - case node.Left != nil && node.Left.Value.SpanOfChildren.ContainsFn(fn): + case node.Left != nil && node.Left.Value.ChildSpan.ContainsFn(fn): node = node.Left - case node.Right != nil && node.Right.Value.SpanOfChildren.ContainsFn(fn): + case node.Right != nil && node.Right.Value.ChildSpan.ContainsFn(fn): node = node.Right default: node = nil @@ -119,24 +114,33 @@ func (t *IntervalTree[K, V]) Search(fn func(K) int) (V, bool) { return zero, false } -func (t *IntervalTree[K, V]) searchAll(fn func(K) int, node *RBNode[intervalValue[K, V]], ret *[]V) { +func (t *IntervalTree[K, V]) Range(fn func(V) bool) { + t.inner.Range(func(node *RBNode[intervalValue[K, V]]) bool { + return fn(node.Value.Val) + }) +} + +func (t *IntervalTree[K, V]) Subrange(rangeFn func(K) int, handleFn func(V) bool) { + t.subrange(t.inner.root, rangeFn, handleFn) +} + +func (t *IntervalTree[K, V]) subrange(node *RBNode[intervalValue[K, V]], rangeFn func(K) int, handleFn func(V) bool) bool { if node == nil { - return + return true } - if !node.Value.SpanOfChildren.ContainsFn(fn) { - return + if !node.Value.ChildSpan.ContainsFn(rangeFn) { + return true } - t.searchAll(fn, node.Left, ret) - if t.keyFn(node.Value).ContainsFn(fn) { - *ret = append(*ret, node.Value.Val) + if !t.subrange(node.Left, rangeFn, handleFn) { + return false } - t.searchAll(fn, node.Right, ret) -} - -func (t *IntervalTree[K, V]) SearchAll(fn func(K) int) []V { - var ret []V - t.searchAll(fn, t.inner.root, &ret) - return ret + if node.Value.ValSpan.ContainsFn(rangeFn) { + if !handleFn(node.Value.Val) { + return false + } + } + if !t.subrange(node.Right, rangeFn, handleFn) { + return false + } + return true } - -// TODO: func (t *IntervalTree[K, V]) Walk(fn func(*RBNode[V]) error) error |