diff --git a/internal/require2/require.go b/internal/require2/require.go index 65c5418..90f3431 100644 --- a/internal/require2/require.go +++ b/internal/require2/require.go @@ -152,6 +152,30 @@ func ElementsMatch[T comparable](t *testing.T, a []T, b []T) { } } +func ElementsEqual[T comparable](t *testing.T, a []T, b []T) { + aSet := make(map[T]int, len(a)) + for _, ai := range a { + aSet[ai] += 1 + } + bSet := make(map[T]int, len(b)) + for _, bi := range b { + bSet[bi] += 1 + } + + for elem, aCount := range aSet { + bCount, _ := bSet[elem] + if aCount != bCount { + fatalf(t, "%#v appears %d times in a and %d times in b", elem, aCount, bCount) + } + } + for elem, bCount := range bSet { + aCount, _ := aSet[elem] + if bCount != aCount { + fatalf(t, "%#v appears %d times in b and %d times in a", elem, bCount, aCount) + } + } +} + func fatalf(t *testing.T, s string, fmtArgs ...any) { var buf [64]uintptr var ptrs []uintptr diff --git a/xslices/xslices.go b/xslices/xslices.go index 63e918d..6950c82 100644 --- a/xslices/xslices.go +++ b/xslices/xslices.go @@ -407,3 +407,36 @@ func uniqueInto[T comparable](into []T, s []T) []T { } return into } + +// Intersect returns a slice containing elements common to each in. If an element appear at least X +// times in each in, it will appear X times in the returned slice. Intersect makes no guarantees +// about the order of the returned elements. +// +// Intersect returns a nil slice when ins have no elements in common. +func Intersect[T comparable](in ...[]T) []T { + visits := map[T][]int{} + for i, s := range in { + for _, e := range s { + if visits[e] == nil { + visits[e] = make([]int, len(in)) + } + visits[e][i] += 1 + } + } + + var out []T + for k, v := range visits { + min := Reduce(v, -1, func(a int, e int) int { + if a < 0 || a > e { + return e + } + return a + }) + + for i := 0; i < min; i++ { + out = append(out, k) + } + } + + return out +} diff --git a/xslices/xslices_example_test.go b/xslices/xslices_example_test.go index d459b50..5aa990a 100644 --- a/xslices/xslices_example_test.go +++ b/xslices/xslices_example_test.go @@ -440,3 +440,16 @@ func ExampleUniqueInPlace() { // Output: // [a b c] } + +func ExampleIntersect() { + intersection := xslices.Intersect( + []string{"a", "b", "b", "c", "d"}, + []string{"a", "b", "b", "c"}, + []string{"a", "b", "b"}, + ) + + fmt.Println(intersection) + + // Output: + // [a b b] +} diff --git a/xslices/xslices_test.go b/xslices/xslices_test.go index 08c04e8..73dfeec 100644 --- a/xslices/xslices_test.go +++ b/xslices/xslices_test.go @@ -51,3 +51,27 @@ func FuzzRemoveUnordered(f *testing.F) { require2.ElementsMatch(t, expected, actual) }) } + +func TestIntersect(t *testing.T) { + testCases := map[string]struct { + in [][]int + want []int + }{ + "nil": {}, + "empty 1": {in: [][]int{}}, + "empty 2": {in: [][]int{{}}}, + "empty 3": {in: [][]int{{}, {}, {}}}, + "empty inter 1": {in: [][]int{{1, 2}, {1, 2}, {}}}, + "empty inter 2": {in: [][]int{{1, 2}, {3, 4}, {5, 6}}}, + "empty inter 3": {in: [][]int{{1, 1}, {2, 2}}}, + "single inter": {in: [][]int{{1, 2, 3}, {1, 2}, {2, 3}}, want: []int{2}}, + "multiple inter": {in: [][]int{{1, 2, 3}, {1, 2}, {1, 2}}, want: []int{1, 2}}, + "complete inter": {in: [][]int{{1, 2, 3}, {3, 1, 2}, {2, 3, 1}}, want: []int{1, 2, 3}}, + "repeated inter": {in: [][]int{{1, 1, 2}, {1, 1, 1}, {1, 1, 3}}, want: []int{1, 1}}, + } + for name, tc := range testCases { + t.Run(name, func(t *testing.T) { + require2.ElementsEqual(t, tc.want, Intersect(tc.in...)) + }) + } +}