diff options
Diffstat (limited to 'lib/containers')
-rw-r--r-- | lib/containers/rbtree.go | 7 | ||||
-rw-r--r-- | lib/containers/rbtree_test.go | 1 |
2 files changed, 8 insertions, 0 deletions
diff --git a/lib/containers/rbtree.go b/lib/containers/rbtree.go index ab79a26..e2d8f8a 100644 --- a/lib/containers/rbtree.go +++ b/lib/containers/rbtree.go @@ -37,6 +37,11 @@ type RBTree[K Ordered[K], V any] struct { KeyFn func(V) K AttrFn func(*RBNode[V]) root *RBNode[V] + len int +} + +func (t *RBTree[K, V]) Len() int { + return t.len } func (t *RBTree[K, V]) Walk(fn func(*RBNode[V]) error) error { @@ -324,6 +329,7 @@ func (t *RBTree[K, V]) Insert(val V) { exact.Value = val return } + t.len++ node := &RBNode[V]{ Color: Red, @@ -394,6 +400,7 @@ func (t *RBTree[K, V]) Delete(key K) { if nodeToDelete == nil { return } + t.len-- // This is closely based on the algorithm presented in CLRS // 3e. diff --git a/lib/containers/rbtree_test.go b/lib/containers/rbtree_test.go index 9841d26..a487c52 100644 --- a/lib/containers/rbtree_test.go +++ b/lib/containers/rbtree_test.go @@ -112,6 +112,7 @@ func checkRBTree[K constraints.Ordered, V any](t *testing.T, expectedSet Set[K], return nil })) require.Equal(t, expectedOrder, actOrder) + require.Equal(t, tree.Len(), len(expectedSet)) } func FuzzRBTree(f *testing.F) { |