diff options
author | Luke Shumaker <lukeshu@lukeshu.com> | 2022-12-20 21:33:28 -0700 |
---|---|---|
committer | Luke Shumaker <lukeshu@lukeshu.com> | 2022-12-21 13:54:10 -0700 |
commit | aa923f4a8bdd679eb2e856780c795f93c4dc17b2 (patch) | |
tree | a85f46853119232a2bececa474a6e0530fb68bc2 | |
parent | 3cd6f42caa633b6726b7b4aebeef92c86b2cf213 (diff) |
containers: sortedmap: Add new .Search() and .SearchAll() methods
-rw-r--r-- | lib/containers/sortedmap.go | 30 |
1 files changed, 24 insertions, 6 deletions
diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go index b759fa3..d6d2f9e 100644 --- a/lib/containers/sortedmap.go +++ b/lib/containers/sortedmap.go @@ -8,13 +8,13 @@ import ( "errors" ) -type orderedKV[K Ordered[K], V any] struct { +type OrderedKV[K Ordered[K], V any] struct { K K V V } type SortedMap[K Ordered[K], V any] struct { - inner RBTree[K, orderedKV[K, V]] + inner RBTree[K, OrderedKV[K, V]] } func (m *SortedMap[K, V]) init() { @@ -23,7 +23,7 @@ func (m *SortedMap[K, V]) init() { } } -func (m *SortedMap[K, V]) keyFn(kv orderedKV[K, V]) K { +func (m *SortedMap[K, V]) keyFn(kv OrderedKV[K, V]) K { return kv.K } @@ -46,7 +46,7 @@ var errStop = errors.New("stop") func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) { m.init() - _ = m.inner.Walk(func(node *RBNode[orderedKV[K, V]]) error { + _ = m.inner.Walk(func(node *RBNode[OrderedKV[K, V]]) error { if f(node.Value.K, node.Value.V) { return nil } else { @@ -57,7 +57,7 @@ func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) { func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) bool) { m.init() - kvs := m.inner.SearchRange(func(kv orderedKV[K, V]) int { + kvs := m.inner.SearchRange(func(kv OrderedKV[K, V]) int { return rangeFn(kv.K, kv.V) }) for _, kv := range kvs { @@ -69,8 +69,26 @@ func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) b func (m *SortedMap[K, V]) Store(key K, value V) { m.init() - m.inner.Insert(orderedKV[K, V]{ + m.inner.Insert(OrderedKV[K, V]{ K: key, V: value, }) } + +func (m *SortedMap[K, V]) Search(fn func(K, V) int) (K, V, bool) { + node := m.inner.Search(func(kv OrderedKV[K, V]) int { + return fn(kv.K, kv.V) + }) + if node == nil { + var zeroK K + var zeroV V + return zeroK, zeroV, false + } + return node.Value.K, node.Value.V, true +} + +func (m *SortedMap[K, V]) SearchAll(fn func(K, V) int) []OrderedKV[K, V] { + return m.inner.SearchRange(func(kv OrderedKV[K, V]) int { + return fn(kv.K, kv.V) + }) +} |