feat: improve Union (#245)

Co-authored-by: hhu <[email protected]>

intersect.go

@@ -141,35 +141,18 @@ func Difference[T comparable](list1 []T, list2 []T) ([]T, []T) {
 	return left, right
 }
 
-// Union returns all distinct elements from both collections.
+// Union returns all distinct elements from given collections.
 // result returns will not change the order of elements relatively.
-func Union[T comparable](list1 []T, list2 []T) []T {
+func Union[T comparable](lists ...[]T) []T {
 	result := []T{}
-
 	seen := map[T]struct{}{}
-	hasAdd := map[T]struct{}{}
-
-	for _, e := range list1 {
-		seen[e] = struct{}{}
-	}
 
-	for _, e := range list2 {
-		seen[e] = struct{}{}
-	}
-
-	for _, e := range list1 {
-		if _, ok := seen[e]; ok {
-			result = append(result, e)
-			hasAdd[e] = struct{}{}
-		}
-	}
-
-	for _, e := range list2 {
-		if _, ok := hasAdd[e]; ok {
-			continue
-		}
-		if _, ok := seen[e]; ok {
-			result = append(result, e)
+	for _, list := range lists {
+		for _, e := range list {
+			if _, ok := seen[e]; !ok {
+				seen[e] = struct{}{}
+				result = append(result, e)
+			}
 		}
 	}
 

intersect_test.go

@@ -220,6 +220,17 @@ func TestUnion(t *testing.T) {
 	is.Equal(result3, []int{0, 1, 2, 3, 4, 5})
 	is.Equal(result4, []int{0, 1, 2})
 	is.Equal(result5, []int{})
+
+	result11 := Union([]int{0, 1, 2, 3, 4, 5}, []int{0, 2, 10}, []int{0, 1, 11})
+	result12 := Union([]int{0, 1, 2, 3, 4, 5}, []int{6, 7}, []int{8, 9})
+	result13 := Union([]int{0, 1, 2, 3, 4, 5}, []int{}, []int{})
+	result14 := Union([]int{0, 1, 2}, []int{0, 1, 2}, []int{0, 1, 2})
+	result15 := Union([]int{}, []int{}, []int{})
+	is.Equal(result11, []int{0, 1, 2, 3, 4, 5, 10, 11})
+	is.Equal(result12, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
+	is.Equal(result13, []int{0, 1, 2, 3, 4, 5})
+	is.Equal(result14, []int{0, 1, 2})
+	is.Equal(result15, []int{})
 }
 
 func TestWithout(t *testing.T) {