update

internal/iterator/iterator.go

@@ -0,0 +1,32 @@
+// Copyright 2021 ecodeclub
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iterator
+
+// 一个迭代器的接口，所有的容器类型都可以实现自己的迭代器
+// 只需要继承当前接口即可
+type Iterator[T any] interface {
+	// 迭代器移动到下一个节点
+	// 如果没有下一个节点，则迭代器所指向的位置变为非法位置，一般情况下为nil
+	Next()
+
+	// 获取迭代器当前所指向的节点的信息
+	Get() (T, error)
+
+	// 判断是否有后继节点
+	HasNext() bool
+
+	// 判断当前节点是否合法
+	Valid() bool
+}

internal/tree/red_black_tree.go

@@ -18,6 +18,8 @@ import (
 	"errors"
 
 	"github.com/ecodeclub/ekit"
+	"github.com/ecodeclub/ekit/internal/iterator"
+	"github.com/ecodeclub/ekit/tuple/pair"
 )
 
 type color bool
@@ -31,6 +33,9 @@ var (
 	ErrRBTreeSameRBNode = errors.New("ekit: RBTree不能添加重复节点Key")
 	ErrRBTreeNotRBNode  = errors.New("ekit: RBTree不存在节点Key")
 	// errRBTreeCantRepaceNil = errors.New("ekit: RBTree不能将节点替换为nil")
+	ErrRBTreeIteratorNoNext  = errors.New("ekit: RBTree Iterator没有后继节点")
+	ErrRBTreeIteratorNodeNil = errors.New("ekit: RBTree Iterator指向的节点为nil")
+	ErrRBTreeIteratorInvalid = errors.New("ekit: RBTree Iterator不合法")
 )
 
 type RBTree[K any, V any] struct {
@@ -53,6 +58,26 @@ type rbNode[K any, V any] struct {
 	left, right, parent *rbNode[K, V]
 }
 
+func (node *rbNode[K, V]) getNext() *rbNode[K, V] {
+	if node == nil {
+		return nil
+	} else if node.right != nil {
+		p := node.right
+		for p.left != nil {
+			p = p.left
+		}
+		return p
+	} else {
+		p := node.parent
+		ch := node
+		for p != nil && ch == p.right {
+			ch = p
+			p = p.parent
+		}
+		return p
+	}
+}
+
 func (node *rbNode[K, V]) setNode(v V) {
 	if node == nil {
 		return
@@ -79,6 +104,15 @@ func newRBNode[K any, V any](key K, value V) *rbNode[K, V] {
 	}
 }
 
+// 获取起始点的迭代器
+func (rb *RBTree[K, V]) Begin() iterator.Iterator[pair.Pair[K, V]] {
+	curr := rb.root
+	for curr.left != nil {
+		curr = curr.left
+	}
+	return newRBTreeIterator(curr)
+}
+
 // Add 增加节点
 func (rb *RBTree[K, V]) Add(key K, value V) error {
 	return rb.addNode(newRBNode(key, value))
@@ -95,6 +129,17 @@ func (rb *RBTree[K, V]) Delete(key K) (V, bool) {
 	return v, false
 }
 
+// 删除节点，但是是使用iterator来删除
+func (rb *RBTree[K, V]) DeleteIt(it iterator.Iterator[pair.Pair[K, V]]) (err error) {
+	iter := it.(*rbTreeIterator[K, V])
+	if iter.node != nil {
+		rb.deleteNode(iter.node)
+		return
+	} else {
+		return ErrRBTreeIteratorInvalid
+	}
+}
+
 // Find 查找节点
 func (rb *RBTree[K, V]) Find(key K) (V, error) {
 	var v V
@@ -103,6 +148,16 @@ func (rb *RBTree[K, V]) Find(key K) (V, error) {
 	}
 	return v, ErrRBTreeNotRBNode
 }
+
+// 查找结点 （但是返回iterator）
+func (rb *RBTree[K, V]) FindIt(key K) (iterator.Iterator[pair.Pair[K, V]], error) {
+	if node := rb.findNode(key); node != nil {
+		return newRBTreeIterator(node), nil
+	}
+	return nil, ErrRBTreeNotRBNode
+}
+
+// 给对应的Key 设置 Value
 func (rb *RBTree[K, V]) Set(key K, value V) error {
 	if node := rb.findNode(key); node != nil {
 		node.setNode(value)
@@ -243,24 +298,7 @@ func (rb *RBTree[K, V]) deleteNode(tgt *rbNode[K, V]) {
 // case1: node节点存在右子节点,则右子树的最小节点是node的后继节点
 // case2: node节点不存在右子节点,则其第一个为左节点的祖先的父节点为node的后继节点
 func (rb *RBTree[K, V]) findSuccessor(node *rbNode[K, V]) *rbNode[K, V] {
-	if node == nil {
-		return nil
-	} else if node.right != nil {
-		p := node.right
-		for p.left != nil {
-			p = p.left
-		}
-		return p
-	} else {
-		p := node.parent
-		ch := node
-		for p != nil && ch == p.right {
-			ch = p
-			p = p.parent
-		}
-		return p
-	}
-
+	return node.getNext()
 }
 
 func (rb *RBTree[K, V]) findNode(key K) *rbNode[K, V] {
@@ -544,3 +582,34 @@ func (node *rbNode[K, V]) getBrother() *rbNode[K, V] {
 	}
 	return node.getParent().getLeft()
 }
+
+type rbTreeIterator[K any, V any] struct {
+	node *rbNode[K, V]
+}
+
+func (iter *rbTreeIterator[K, V]) Next() {
+	iter.node = iter.node.getNext()
+}
+
+func (iter *rbTreeIterator[K, V]) HasNext() bool {
+	return iter.node.getNext() != nil
+}
+
+func (iter *rbTreeIterator[K, V]) Get() (kvPair pair.Pair[K, V], err error) {
+	if iter.node == nil {
+		err = ErrRBTreeIteratorInvalid
+		return
+	}
+	kvPair = pair.NewPair(iter.node.key, iter.node.value)
+	return
+}
+
+func (iter *rbTreeIterator[K, V]) Valid() bool {
+	return iter.node != nil
+}
+
+func newRBTreeIterator[K any, V any](node *rbNode[K, V]) iterator.Iterator[pair.Pair[K, V]] {
+	return &rbTreeIterator[K, V]{
+		node: node,
+	}
+}

internal/tree/red_black_tree_iterator_test.go

@@ -0,0 +1,40 @@
+package tree
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/ecodeclub/ekit"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIteratorToVisitFullRBTree(t *testing.T) {
+	n := 10000
+	arr := generateArray(n)
+	rbTree := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+	for _, v := range arr {
+		assert.Nil(t, rbTree.Add(v, v))
+	}
+
+	arrVisit := make([]int, n)
+	id := 0
+	for iter := rbTree.Begin(); iter.Valid(); iter.Next() {
+		pa, err := iter.Get()
+		assert.Nil(t, err)
+		arrVisit[id] = pa.Key
+		assert.Equal(t, id, pa.Key)
+		id++
+	}
+	assert.Equal(t, n, id)
+}
+
+func generateArray(n int) []int {
+	res := make([]int, n)
+	for i := 0; i < n; i++ {
+		res[i] = i
+	}
+	rand.Shuffle(n, func(i, j int) {
+		res[i], res[j] = res[j], res[i]
+	})
+	return res
+}