From 6a2a333ded675666f68ffafa07ddc9a5e093cd32 Mon Sep 17 00:00:00 2001
From: H <hhu.cc@qq.com>
Date: Fri, 14 Oct 2022 23:55:11 +0800
Subject: [PATCH] feat: improve Union (#245)

Co-authored-by: hhu <hhu@tencent.com>
---
 intersect.go      | 33 ++++++++-------------------------
 intersect_test.go | 11 +++++++++++
 2 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/intersect.go b/intersect.go
index 38a0052b..cf6cab3d 100644
--- a/intersect.go
+++ b/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)
+			}
 		}
 	}
 
diff --git a/intersect_test.go b/intersect_test.go
index 8867d8f9..339dbcbb 100644
--- a/intersect_test.go
+++ b/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) {