Repository: ecodeclub/ekit
Pull request #278: 增加treemap的Iterate
Author: login: dxyinme
Created: 2025-05-09T14:15:03Z | Merged: 2025-05-10T02:31:53Z



Diff:
diff --git a/internal/tree/red_black_tree.go b/internal/tree/red_black_tree.go
index cec0eac0..f07883f5 100644
--- a/internal/tree/red_black_tree.go
+++ b/internal/tree/red_black_tree.go
@@ -118,16 +118,24 @@ func (rb *RBTree[K, V]) KeyValues() ([]K, []V) {
 	if rb.root == nil {
 		return keys, values
 	}
-	rb.inOrderTraversal(func(node *rbNode[K, V]) {
+	rb.inOrderTraversal(func(node *rbNode[K, V]) bool {
 		keys = append(keys, node.key)
 		values = append(values, node.value)
+		return true
 	})
 	return keys, values
 }
 
+// Iterate 按照key的顺序遍历并执行cb，如果cb返回值为false则结束遍历，否则继续遍历
+func (rb *RBTree[K, V]) Iterate(cb func(key K, value V) bool) {
+	rb.inOrderTraversal(func(node *rbNode[K, V]) bool {
+		return cb(node.key, node.value)
+	})
+}
+
 // inOrderTraversal 中序遍历
-func (rb *RBTree[K, V]) inOrderTraversal(visit func(node *rbNode[K, V])) {
-	stack := make([]*rbNode[K, V], 0, rb.size)
+func (rb *RBTree[K, V]) inOrderTraversal(visit func(node *rbNode[K, V]) bool) {
+	stack := make([]*rbNode[K, V], 0)
 	curr := rb.root
 	for curr != nil || len(stack) > 0 {
 		for curr != nil {
@@ -136,7 +144,9 @@ func (rb *RBTree[K, V]) inOrderTraversal(visit func(node *rbNode[K, V])) {
 		}
 		curr = stack[len(stack)-1]
 		stack = stack[:len(stack)-1]
-		visit(curr)
+		if !visit(curr) {
+			break
+		}
 		curr = curr.right
 	}
 }
diff --git a/internal/tree/red_black_tree_test.go b/internal/tree/red_black_tree_test.go
index 494217dd..a945f313 100644
--- a/internal/tree/red_black_tree_test.go
+++ b/internal/tree/red_black_tree_test.go
@@ -1397,6 +1397,73 @@ func TestRBTree_KeyValues(t *testing.T) {
 	}
 }
 
+func TestRBTree_Iterate(t *testing.T) {
+	for _, testCase := range []struct {
+		name        string
+		expectedLen int
+		inputStart  int
+		inputEnd    int
+		// 如果为true则遍历结束
+		endConditionFunc func(key int) bool
+	}{
+		{
+			name:        "treeMap为空",
+			expectedLen: 0,
+			inputStart:  1,
+			inputEnd:    0,
+			endConditionFunc: func(key int) bool {
+				return false
+			},
+		},
+		{
+			name:        "treeMap 有10000个元素,遍历所有小于等于8000的元素",
+			expectedLen: 8000,
+			inputStart:  1,
+			inputEnd:    10000,
+			endConditionFunc: func(key int) bool {
+				return key > 8000
+			},
+		},
+		{
+			name:        "treeMap 有10000个元素,遍历所有元素",
+			expectedLen: 10000,
+			inputStart:  1,
+			inputEnd:    10000,
+			endConditionFunc: func(key int) bool {
+				return false
+			},
+		},
+		{
+			name:        "treeMap 有10个元素,由于第一个就不符合条件所以遍历立刻中断",
+			expectedLen: 0,
+			inputStart:  1,
+			inputEnd:    10,
+			endConditionFunc: func(key int) bool {
+				return key < 5
+			},
+		},
+	} {
+		t.Run(testCase.name, func(t *testing.T) {
+			rbTree := NewRBTree[int, int](compare())
+			for i := testCase.inputStart; i <= testCase.inputEnd; i++ {
+				assert.Nil(t, rbTree.Add(i, i))
+			}
+			arr := make([]int, 0)
+			rbTree.Iterate(func(key, value int) bool {
+				if testCase.endConditionFunc(key) {
+					return false
+				}
+				arr = append(arr, value)
+				return true
+			})
+			assert.Equal(t, testCase.expectedLen, len(arr))
+			for i := 0; i < testCase.expectedLen; i++ {
+				assert.Equal(t, testCase.inputStart+i, arr[i])
+			}
+		})
+	}
+}
+
 // IsRedBlackTree 检测是否满足红黑树
 func IsRedBlackTree[K any, V any](root *rbNode[K, V]) bool {
 	// 检测节点是否黑色
diff --git a/mapx/treemap.go b/mapx/treemap.go
index 40f611fa..dabdbfc3 100644
--- a/mapx/treemap.go
+++ b/mapx/treemap.go
@@ -101,4 +101,9 @@ func (treeMap *TreeMap[T, V]) Len() int64 {
 	return int64(treeMap.tree.Size())
 }
 
+// Iterate 按照key的顺序遍历并执行cb，如果cb返回值为false则结束遍历，否则继续遍历
+func (treeMap *TreeMap[K, V]) Iterate(cb func(key K, value V) bool) {
+	treeMap.tree.Iterate(cb)
+}
+
 var _ mapi[any, any] = (*TreeMap[any, any])(nil)
diff --git a/mapx/treemap_example_test.go b/mapx/treemap_example_test.go
index b00fe803..02f871f1 100644
--- a/mapx/treemap_example_test.go
+++ b/mapx/treemap_example_test.go
@@ -29,3 +29,26 @@ func ExampleNewTreeMap() {
 	// Output:
 	// 11
 }
+
+func ExampleTreeMap_Iterate() {
+	m, _ := mapx.NewTreeMap[int, int](ekit.ComparatorRealNumber[int])
+	_ = m.Put(1, 11)
+	_ = m.Put(-1, 12)
+	_ = m.Put(100, 13)
+	_ = m.Put(-100, 14)
+	_ = m.Put(-101, 15)
+
+	m.Iterate(func(key, value int) bool {
+		if key > 1 {
+			return false
+		}
+		fmt.Println(key, value)
+		return true
+	})
+
+	// Output:
+	// -101 15
+	// -100 14
+	// -1 12
+	// 1 11
+}
diff --git a/mapx/treemap_test.go b/mapx/treemap_test.go
index c224631a..a237555c 100644
--- a/mapx/treemap_test.go
+++ b/mapx/treemap_test.go
@@ -445,6 +445,74 @@ func TestTreeMap_Len(t *testing.T) {
 	}
 }
 
+func TestRBTree_Iterate(t *testing.T) {
+	for _, testCase := range []struct {
+		name        string
+		expectedLen int
+		inputStart  int
+		inputEnd    int
+		// 如果为true则遍历结束
+		endConditionFunc func(key int) bool
+	}{
+		{
+			name:        "treeMap为空",
+			expectedLen: 0,
+			inputStart:  1,
+			inputEnd:    0,
+			endConditionFunc: func(key int) bool {
+				return false
+			},
+		},
+		{
+			name:        "treeMap 有10000个元素,遍历所有小于等于8000的元素",
+			expectedLen: 8000,
+			inputStart:  1,
+			inputEnd:    10000,
+			endConditionFunc: func(key int) bool {
+				return key > 8000
+			},
+		},
+		{
+			name:        "treeMap 有10000个元素,遍历所有元素",
+			expectedLen: 10000,
+			inputStart:  1,
+			inputEnd:    10000,
+			endConditionFunc: func(key int) bool {
+				return false
+			},
+		},
+		{
+			name:        "treeMap 有10个元素,由于第一个就不符合条件所以遍历立刻中断",
+			expectedLen: 0,
+			inputStart:  1,
+			inputEnd:    10,
+			endConditionFunc: func(key int) bool {
+				return key < 5
+			},
+		},
+	} {
+		t.Run(testCase.name, func(t *testing.T) {
+			treeMap, err := NewTreeMap[int, int](compare())
+			assert.Nil(t, err)
+			for i := testCase.inputStart; i <= testCase.inputEnd; i++ {
+				assert.Nil(t, treeMap.Put(i, i))
+			}
+			arr := make([]int, 0)
+			treeMap.Iterate(func(key, value int) bool {
+				if testCase.endConditionFunc(key) {
+					return false
+				}
+				arr = append(arr, value)
+				return true
+			})
+			assert.Equal(t, testCase.expectedLen, len(arr))
+			for i := 0; i < testCase.expectedLen; i++ {
+				assert.Equal(t, testCase.inputStart+i, arr[i])
+			}
+		})
+	}
+}
+
 func compare() ekit.Comparator[int] {
 	return ekit.ComparatorRealNumber[int]
 }


Reviews:
user: login: flycash
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T14:39:28Z
user: login: dxyinme
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-10T02:10:30Z
user: login: flycash
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-10T02:31:35Z

Review comments:
diff_hunk: @@ -1397,6 +1397,29 @@ func TestRBTree_KeyValues(t *testing.T) {
 	}
 }
 
+// 测试遍历
+// 插入10000个数字，遍历到第8000个为止，并且是按顺序遍历
+func TestRBTree_Iterate(t *testing.T) {
path: internal/tree/red_black_tree_test.go
user: login: flycash
body: 用 Table driven 的形式来写。提供：
- 里面一个元素都没有
- 遍历一部分元素之后，中断；
- 遍历完全部
created_at: 2025-05-09T14:39:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1397,6 +1397,29 @@ func TestRBTree_KeyValues(t *testing.T) {
 	}
 }
 
+// 测试遍历
+// 插入10000个数字，遍历到第8000个为止，并且是按顺序遍历
+func TestRBTree_Iterate(t *testing.T) {
path: internal/tree/red_black_tree_test.go
user: login: dxyinme
body: 已修改
created_at: 2025-05-10T02:10:30Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-09T14:20:32Z
body: ## [Codecov](https://app.codecov.io/gh/ecodeclub/ekit/pull/278?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecodeclub) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.59%. Comparing base [(`40e0aa7`)](https://app.codecov.io/gh/ecodeclub/ekit/commit/40e0aa7f70b6a80af657e9ef54939380c70b7dee?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecodeclub) to head [(`4e15b0a`)](https://app.codecov.io/gh/ecodeclub/ekit/commit/4e15b0acb5b2bfff18d849001e165f10ec6a50ef?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecodeclub).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev     #278      +/-   ##
==========================================
- Coverage   94.69%   94.59%   -0.10%     
==========================================
  Files          75       75              
  Lines        3767     3774       +7     
==========================================
+ Hits         3567     3570       +3     
- Misses        161      165       +4     
  Partials       39       39              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecodeclub/ekit/pull/278?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecodeclub).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecodeclub).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: flycash
created_at: 2025-05-10T02:31:48Z
body: 可以把其它 map 的遍历也加一下了。
author_association: CONTRIBUTOR
user: login: dxyinme
created_at: 2025-05-10T02:39:14Z
body: > 可以把其它 map 的遍历也加一下了。

好的，我过两天有空了搞一下
author_association: CONTRIBUTOR