迭代器设计

internal/iterator/iterator.go

@@ -0,0 +1,38 @@
+// 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 {
+	// 迭代器移动到下一个节点，如果成功的话就返回true
+	// 如果没有下一个节点，则迭代器所指向的位置会变为非法，一般为nil，并且返回false
+	Next() bool
+
+	// 获取迭代器当前所指向的节点的信息
+	Get() T
+
+	// 获取error
+	Err() error
+
+	// 判断是否有后继节点
+	HasNext() bool
+
+	// 判断当前节点是否合法
+	Valid() bool
+
+	// 删除当前节点
+	Delete()
+}

internal/tree/red_black_tree.go

@@ -16,8 +16,11 @@ package tree
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/ecodeclub/ekit"
+	"github.com/ecodeclub/ekit/internal/iterator"
+	"github.com/ecodeclub/ekit/tuple/pair"
 )
 
 type color bool
@@ -31,6 +34,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 +59,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 +105,19 @@ func newRBNode[K any, V any](key K, value V) *rbNode[K, V] {
 	}
 }
 
+// 获取起始点的迭代器
+func (rb *RBTree[K, V]) beginNode() *rbNode[K, V] {
+	curr := rb.root
+	for curr.left != nil {
+		curr = curr.left
+	}
+	return curr
+}
+
+func (rb *RBTree[K, V]) Begin() iterator.Iterator[pair.Pair[K, V]] {
+	return newRBTreeIterator(rb, rb.beginNode())
+}
+
 // Add 增加节点
 func (rb *RBTree[K, V]) Add(key K, value V) error {
 	return rb.addNode(newRBNode(key, value))
@@ -95,6 +134,14 @@ func (rb *RBTree[K, V]) Delete(key K) (V, bool) {
 	return v, false
 }
 
+// 查找结点 （但是返回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(rb, node), nil
+	}
+	return nil, ErrRBTreeNotRBNode
+}
+
 // Find 查找节点
 func (rb *RBTree[K, V]) Find(key K) (V, error) {
 	var v V
@@ -103,6 +150,8 @@ func (rb *RBTree[K, V]) Find(key K) (V, error) {
 	}
 	return v, 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 +292,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 +576,88 @@ func (node *rbNode[K, V]) getBrother() *rbNode[K, V] {
 	}
 	return node.getParent().getLeft()
 }
+
+func (node *rbNode[K, V]) isValidNode() bool {
+	if node == nil {
+		return false
+	}
+	vis := 0
+	if node.getLeft() != nil {
+		if node.getLeft().getParent() != node {
+			return false
+		}
+		vis++
+	}
+	if node.getRight() != nil {
+		if node.getRight().getParent() != node {
+			return false
+		}
+		vis++
+	}
+	if node.getParent() != nil {
+		if node.getParent().getLeft() != node && node.getParent().getRight() != node {
+			return false
+		}
+		vis++
+	}
+	return vis > 0
+}
+
+type rbTreeIterator[K any, V any] struct {
+	rbTree   *RBTree[K, V]
+	currNode *rbNode[K, V]
+	nxtNode  *rbNode[K, V]
+	err      error
+}
+
+func (iter *rbTreeIterator[K, V]) Next() bool {
+	fmt.Println(iter.currNode.key, iter.currNode.value)
+	iter.err = nil
+	iter.currNode = iter.nxtNode
+	if iter.currNode != nil {
+		iter.nxtNode = iter.currNode.getNext()
+		return true
+	}
+	iter.err = ErrRBTreeIteratorNoNext
+	return false
+}
+
+func (iter *rbTreeIterator[K, V]) HasNext() bool {
+	return iter.nxtNode != nil
+}
+
+func (iter *rbTreeIterator[K, V]) Get() (kvPair pair.Pair[K, V]) {
+	iter.err = nil
+	if iter.currNode == nil {
+		iter.err = ErrRBTreeIteratorInvalid
+		return
+	}
+	kvPair = pair.NewPair(iter.currNode.key, iter.currNode.value)
+	return
+}
+
+func (iter *rbTreeIterator[K, V]) Err() error {
+	return iter.err
+}
+
+func (iter *rbTreeIterator[K, V]) Valid() bool {
+	return iter.currNode != nil
+}
+
+func (iter *rbTreeIterator[K, V]) Delete() {
+	iter.err = nil
+	if !iter.currNode.isValidNode() {
+		iter.err = ErrRBTreeIteratorInvalid
+		return
+	}
+	iter.rbTree.deleteNode(iter.currNode)
+}
+
+func newRBTreeIterator[K any, V any](rbTree *RBTree[K, V], rbNode *rbNode[K, V]) iterator.Iterator[pair.Pair[K, V]] {
+	iter := &rbTreeIterator[K, V]{
+		rbTree:   rbTree,
+		currNode: rbNode,
+	}
+	iter.nxtNode = iter.currNode.getNext()
+	return iter
+}

internal/tree/red_black_tree_iterator_test.go

@@ -0,0 +1,147 @@
+// 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 tree
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/ecodeclub/ekit"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIteratorToVisitFullRBTree(t *testing.T) {
+	t.Parallel()
+	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(), iter.Err()
+		assert.Nil(t, err)
+		arrVisit[id] = pa.Key
+		assert.Equal(t, id, pa.Key)
+		id++
+	}
+	assert.Equal(t, n, id)
+}
+
+func TestIteratorFind(t *testing.T) {
+	t.Run("查找存在的节点", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		it, err := rbt.FindIt(-100)
+		assert.Nil(t, err)
+		assert.Equal(t, 102, it.Get().Value)
+	})
+
+	t.Run("查找不存在的节点", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		it, err := rbt.FindIt(2)
+		assert.Equal(t, ErrRBTreeNotRBNode, err)
+		assert.Nil(t, it)
+	})
+
+	t.Run("查找存在的节点，删除后不存在", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		it, err := rbt.FindIt(-100)
+		assert.Nil(t, err)
+		assert.Equal(t, 102, it.Get().Value)
+		it.Delete()
+		assert.Nil(t, it.Err())
+		it, err = rbt.FindIt(-100)
+		assert.Equal(t, ErrRBTreeNotRBNode, err)
+		assert.Nil(t, it)
+	})
+
+	t.Run("查找不存在的节点,增加后存在", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		it, err := rbt.FindIt(2)
+		assert.Equal(t, ErrRBTreeNotRBNode, err)
+		assert.Nil(t, it)
+		assert.Nil(t, rbt.Add(2, 104))
+		it, err = rbt.FindIt(2)
+		assert.Nil(t, err)
+		assert.Equal(t, 104, it.Get().Value)
+	})
+}
+
+func TestIteratorDelete(t *testing.T) {
+	t.Run("重复删除某个节点", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		it, err := rbt.FindIt(-100)
+		assert.Nil(t, err)
+		it.Delete()
+		assert.Equal(t, nil, it.Err())
+		it.Delete()
+		assert.Equal(t, ErrRBTreeIteratorInvalid, it.Err())
+	})
+	t.Run("删除节点后正常遍历", func(t *testing.T) {
+		t.Parallel()
+		rbt := NewRBTree[int, int](ekit.ComparatorRealNumber[int])
+		assert.Nil(t, rbt.Add(1, 101))
+		assert.Nil(t, rbt.Add(-100, 102))
+		assert.Nil(t, rbt.Add(100, 103))
+		assert.Nil(t, rbt.Add(101, 104))
+		assert.Nil(t, rbt.Add(102, 105))
+
+		result := make([]int, 0)
+		for it := rbt.Begin(); it.Valid(); it.Next() {
+			key := it.Get().Key
+			if key == 100 {
+				it.Delete()
+				assert.Nil(t, it.Err())
+				continue
+			}
+			result = append(result, key)
+		}
+		assert.EqualValues(t, []int{-100, 1, 101, 102}, result)
+	})
+}
+
+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
+}