diff --git a/slice/slice.go b/slice/slice.go index 86b93e2..7772f54 100644 --- a/slice/slice.go +++ b/slice/slice.go @@ -594,7 +594,7 @@ func UniqueBy[T comparable](slice []T, iteratee func(item T) T) []T { return Unique(result) } -// Union creates a slice of unique values, in order, from all given slices. using == for equality comparisons. +// Union creates a slice of unique values, in order, from all given slices. func Union[T comparable](slices ...[]T) []T { result := []T{} contain := map[T]struct{}{} @@ -611,6 +611,24 @@ func Union[T comparable](slices ...[]T) []T { return result } +// UnionBy is like Union, what's more it accepts iteratee which is invoked for each element of each slice +func UnionBy[T any, V comparable](predicate func(item T) V, slices ...[]T) []T { + result := []T{} + contain := map[V]struct{}{} + + for _, slice := range slices { + for _, item := range slice { + val := predicate(item) + if _, ok := contain[val]; !ok { + contain[val] = struct{}{} + result = append(result, item) + } + } + } + + return result +} + // Intersection creates a slice of unique values that included by all slices. func Intersection[T comparable](slices ...[]T) []T { if len(slices) == 0 { diff --git a/slice/slice_test.go b/slice/slice_test.go index fc3ad5d..51e70df 100644 --- a/slice/slice_test.go +++ b/slice/slice_test.go @@ -428,6 +428,17 @@ func TestUnion(t *testing.T) { assert.Equal([]int{1, 3, 4, 6}, Union(s1)) } +func TestUnionBy(t *testing.T) { + assert := internal.NewAssert(t, "TestUnionBy") + + testFunc := func(i int) int { + return i / 2 + } + + result := UnionBy(testFunc, []int{0, 1, 2, 3, 4, 5}, []int{0, 2, 10}) + assert.Equal(result, []int{0, 2, 4, 10}) +} + func TestIntersection(t *testing.T) { s1 := []int{1, 2, 2, 3} s2 := []int{1, 2, 3, 4}