fix(treeset): fix RemoveFunc not remove all matched node (#63)

zonewave · shoenig · commit 7741a0f88e53 · 2023-05-13T08:08:01.000-05:00
1. add a filterSlice method that returns the elements of s that satisfy the predicate as a slice, in order.
2. first find remove ids by filterSlice, secondly to remove by removeSlice
diff --git a/examples_treeset_test.go b/examples_treeset_test.go
@@ -157,20 +157,20 @@ func ExampleTreeSet_RemoveSet() {
 }
 
 func ExampleTreeSet_RemoveFunc() {
-	s := TreeSetFrom[int, Compare[int]]([]int{1, 2, 3, 4, 5, 6, 7, 8, 9}, Cmp[int])
+	s := TreeSetFrom[int, Compare[int]](ints(20), Cmp[int])
 
 	fmt.Println(s)
 
 	even := func(i int) bool {
-		return i%2 == 0
+		return i%3 != 0
 	}
 	s.RemoveFunc(even)
 
 	fmt.Println(s)
 
 	// Output:
-	// [1 2 3 4 5 6 7 8 9]
-	// [1 3 5 7 9]
+	// [1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]
+	// [3 6 9 12 15 18]
 }
 
 func ExampleTreeSet_Contains() {
diff --git a/treeset.go b/treeset.go
@@ -154,15 +154,8 @@ func (s *TreeSet[T, C]) RemoveSet(o *TreeSet[T, C]) bool {
 //
 // Return true if s was modified, false otherwise.
 func (s *TreeSet[T, C]) RemoveFunc(f func(T) bool) bool {
-	modified := false
-	remove := func(item T) bool {
-		if f(item) && s.Remove(item) {
-			modified = true
-		}
-		return true
-	}
-	s.ForEach(remove)
-	return modified
+	removeIds := s.FilterSlice(f)
+	return s.RemoveSlice(removeIds)
 }
 
 // Min returns the smallest item in the set.
@@ -356,6 +349,18 @@ func (s *TreeSet[T, C]) Slice() []T {
 	return result
 }
 
+// FilterSlice returns the elements of s that satisfy the predicate f.
+func (s *TreeSet[T, C]) FilterSlice(filter func(T) bool) []T {
+	result := make([]T, 0, s.Size())
+	s.ForEach(func(t T) bool {
+		if filter != nil && filter(t) {
+			result = append(result, t)
+		}
+		return true
+	})
+	return result
+}
+
 // Subset returns whether o is a subset of s.
 func (s *TreeSet[T, C]) Subset(o *TreeSet[T, C]) bool {
 	// try the fast paths

Original file line number	Diff line number	Diff line change
`@@ -157,20 +157,20 @@ func ExampleTreeSet_RemoveSet() {`
`157`	`157`	`}`
`158`	`158`
`159`	`159`	`func ExampleTreeSet_RemoveFunc() {`
`160`		`- s := TreeSetFrom[int, Compare[int]]([]int{1, 2, 3, 4, 5, 6, 7, 8, 9}, Cmp[int])`
	`160`	`+ s := TreeSetFrom[int, Compare[int]](ints(20), Cmp[int])`
`161`	`161`
`162`	`162`	`fmt.Println(s)`
`163`	`163`
`164`	`164`	`even := func(i int) bool {`
`165`		`- return i%2 == 0`
	`165`	`+ return i%3 != 0`
`166`	`166`	`}`
`167`	`167`	`s.RemoveFunc(even)`
`168`	`168`
`169`	`169`	`fmt.Println(s)`
`170`	`170`
`171`	`171`	`// Output:`
`172`		`- // [1 2 3 4 5 6 7 8 9]`
`173`		`- // [1 3 5 7 9]`
	`172`	`+ // [1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]`
	`173`	`+ // [3 6 9 12 15 18]`
`174`	`174`	`}`
`175`	`175`
`176`	`176`	`func ExampleTreeSet_Contains() {`