docs: continue with examples for TreeSet

shoenig · shoenig · commit f96040c46ccc · 2023-05-05T17:16:06.000-05:00
diff --git a/examples_treeset_test.go b/examples_treeset_test.go
@@ -88,7 +88,22 @@ func ExampleTreeSet_InsertSlice() {
 	// [blue green red]
 }
 
-// InsertSet
+func ExampleTreeSet_InsertSet() {
+	s1 := TreeSetFrom[string, Compare[string]]([]string{"red", "green"}, Cmp[string])
+	s2 := TreeSetFrom[string, Compare[string]]([]string{"green", "blue"}, Cmp[string])
+
+	fmt.Println(s1)
+	fmt.Println(s2)
+
+	s1.InsertSet(s2)
+
+	fmt.Println(s1)
+
+	// Output:
+	// [green red]
+	// [blue green]
+	// [blue green red]
+}
 
 func ExampleTreeSet_Remove() {
 	s := TreeSetFrom[string, Compare[string]]([]string{"red", "green", "blue"}, Cmp[string])
@@ -124,9 +139,39 @@ func ExampleTreeSet_RemoveSlice() {
 	// [green]
 }
 
-// RemoveSet
+func ExampleTreeSet_RemoveSet() {
+	s1 := TreeSetFrom[string, Compare[string]]([]string{"a", "b", "c", "d", "e", "f"}, Cmp[string])
+	s2 := TreeSetFrom[string, Compare[string]]([]string{"e", "z", "a"}, Cmp[string])
+
+	fmt.Println(s1)
+	fmt.Println(s2)
+
+	s1.RemoveSet(s2)
+
+	fmt.Println(s1)
+
+	// Output:
+	// [a b c d e f]
+	// [a e z]
+	// [b c d f]
+}
+
+func ExampleTreeSet_RemoveFunc() {
+	s := TreeSetFrom[int, Compare[int]]([]int{1, 2, 3, 4, 5, 6, 7, 8, 9}, Cmp[int])
 
-// RemoveFunc
+	fmt.Println(s)
+
+	even := func(i int) bool {
+		return i%2 == 0
+	}
+	s.RemoveFunc(even)
+
+	fmt.Println(s)
+
+	// Output:
+	// [1 2 3 4 5 6 7 8 9]
+	// [1 3 5 7 9]
+}
 
 func ExampleTreeSet_Contains() {
 	s := TreeSetFrom[string, Compare[string]]([]string{"red", "green", "blue"}, Cmp[string])
diff --git a/set.go b/set.go
@@ -159,7 +159,7 @@ func (s *Set[T]) RemoveSet(o *Set[T]) bool {
 // RemoveFunc will remove each element from s that satisfies condition f.
 //
 // Return true if s was modified, false otherwise.
-func (s *Set[T]) RemoveFunc(f func(item T) bool) bool {
+func (s *Set[T]) RemoveFunc(f func(T) bool) bool {
 	modified := false
 	for item := range s.items {
 		if applies := f(item); applies {
diff --git a/treeset.go b/treeset.go
@@ -101,6 +101,21 @@ func (s *TreeSet[T, C]) InsertSlice(items []T) bool {
 	return modified
 }
 
+// InsertSet will insert each element of o into s.
+//
+// Return true if s was modified (at least one item of o was not already in s), false otherwise.
+func (s *TreeSet[T, C]) InsertSet(o *TreeSet[T, C]) bool {
+	modified := false
+	insert := func(item T) bool {
+		if s.Insert(item) {
+			modified = true
+		}
+		return true
+	}
+	o.ForEach(insert)
+	return modified
+}
+
 // Remove item from s.
 //
 // Returns true if s was modified (item was in s), false otherwise.
@@ -121,6 +136,36 @@ func (s *TreeSet[T, C]) RemoveSlice(items []T) bool {
 	return modified
 }
 
+// RemoveSet will remove each element in o from s.
+//
+// Returns true if s was modified (at least one item in o was in s), false otherwise.
+func (s *TreeSet[T, C]) RemoveSet(o *TreeSet[T, C]) bool {
+	modified := false
+	remove := func(item T) bool {
+		if s.Remove(item) {
+			modified = true
+		}
+		return true
+	}
+	o.ForEach(remove)
+	return modified
+}
+
+// RemoveFunc will remove each element from s that satisifies condition f.
+//
+// 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
+}
+
 // Min returns the smallest item in the set.
 //
 // Must not be called on an empty set.
diff --git a/treeset_test.go b/treeset_test.go
@@ -6,11 +6,12 @@ package set
 import (
 	"context"
 	"fmt"
-	"github.com/shoenig/test/must"
-	"go.uber.org/goleak"
 	"math/rand"
 	"strings"
 	"testing"
+
+	"github.com/shoenig/test/must"
+	"go.uber.org/goleak"
 )
 
 const (
@@ -150,6 +151,17 @@ func TestTreeSet_InsertSlice(t *testing.T) {
 	must.False(t, ts.InsertSlice(numbers))
 }
 
+func TestTreeSet_InsertSet(t *testing.T) {
+	cmp := Cmp[int]
+
+	ts1 := TreeSetFrom[int, Compare[int]]([]int{1, 3, 5, 7, 9}, cmp)
+	ts2 := TreeSetFrom[int, Compare[int]]([]int{1, 2, 3}, cmp)
+
+	must.True(t, ts1.InsertSet(ts2))
+	must.Eq(t, []int{1, 2, 3, 5, 7, 9}, ts1.Slice())
+	must.Eq(t, []int{1, 2, 3}, ts2.Slice())
+}
+
 func TestTreeSet_Remove_int(t *testing.T) {
 	cmp := Cmp[int]
 	ts := NewTreeSet[int, Compare[int]](cmp)
@@ -192,6 +204,40 @@ func TestTreeSet_RemoveSlice(t *testing.T) {
 	must.Empty(t, ts)
 }
 
+func TestTreeSet_RemoveSet(t *testing.T) {
+	cmp := Cmp[int]
+
+	ts1 := NewTreeSet[int, Compare[int]](cmp)
+	ts2 := NewTreeSet[int, Compare[int]](cmp)
+
+	numbers := ints(size)
+	random := shuffle(numbers)
+	ts1.InsertSlice(random)
+
+	random2 := shuffle(numbers[5:])
+	ts2.InsertSlice(random2)
+
+	ts1.RemoveSet(ts2)
+	result := ts1.Slice()
+	must.Eq(t, []int{1, 2, 3, 4, 5}, result)
+}
+
+func TestTreeSet_RemoveFunc(t *testing.T) {
+	cmp := Cmp[byte]
+
+	ts := TreeSetFrom[byte, Compare[byte]]([]byte{
+		'a', 'b', '1', 'c', '2', 'd',
+	}, cmp)
+
+	notAlpha := func(c byte) bool {
+		return c < 'a' || c > 'z'
+	}
+
+	ts.RemoveFunc(notAlpha)
+
+	must.Eq(t, []byte{'a', 'b', 'c', 'd'}, ts.Slice())
+}
+
 func TestTreeSet_Contains(t *testing.T) {
 	t.Run("empty", func(t *testing.T) {
 		ts := NewTreeSet[int, Compare[int]](Cmp[int])
