diff --git a/tree.go b/tree.go index 931f99d..7587838 100644 --- a/tree.go +++ b/tree.go @@ -239,6 +239,24 @@ func (t *Tree[K, V]) Remove(key K) (*Tree[K, V], bool) { return t, false } +// ForEach will call f for each node in this tree, in dfs order. If f returns false, interation stops. +func (t *Tree[K, V]) ForEach(f func(K, V) bool) { + if t == nil { + return + } + t.forEach(f) +} + +func (t *Tree[K, V]) forEach(f func(K, V) bool) bool { + if t.left != nil && !t.left.forEach(f) { + return false + } + if !f(t.key, t.value) { + return false + } + return t.right == nil || t.right.forEach(f) +} + // balance checks if, after Insert or Remove, the tree has become unbalanced. // If it has, balance rotates (or double rotates) to fix it, updating childSize as needed. // rightHeavy should be true if a new node was inserted into right (or removed from left), else false. diff --git a/tree_test.go b/tree_test.go index 2a27eb3..3474d16 100644 --- a/tree_test.go +++ b/tree_test.go @@ -159,3 +159,27 @@ func TestTree2Remove(t *testing.T) { t.Fatal("should have nothing left") } } + +func TestTree2ForEach(t *testing.T) { + rand.Seed(time.Now().UnixNano()) + var test *Tree[cmpUint, cmpUint] + for i := 0; i < cmpUintTestSize; i++ { + newVal := cmpUint(rand.Uint64()) + test, _ = test.Insert(newVal, newVal) + } + + keys := test.Keys() + var diy []cmpUint + test.ForEach(func(key, _ cmpUint) bool { + diy = append(diy, key) + return true + }) + if len(diy) != len(keys) { + t.Fatalf("should be same length") + } + for i, key := range keys { + if diy[i] != key { + t.Fatalf("should be equal @%d %v %v", i, key, diy[i]) + } + } +}