diff --git a/util/slicesx/slicesx.go b/util/slicesx/slicesx.go index ce55594db..ba5be7271 100644 --- a/util/slicesx/slicesx.go +++ b/util/slicesx/slicesx.go @@ -42,3 +42,18 @@ func Shuffle[S ~[]T, T any](s S) { s[i], s[j] = s[j], s[i] } } + +// Partition returns two slices, the first containing the elements of the input +// slice for which the callback evaluates to true, the second containing the rest. +// +// This function does not mutate s. +func Partition[S ~[]T, T any](s S, cb func(T) bool) (trues, falses S) { + for _, elem := range s { + if cb(elem) { + trues = append(trues, elem) + } else { + falses = append(falses, elem) + } + } + return +} diff --git a/util/slicesx/slicesx_test.go b/util/slicesx/slicesx_test.go index 1d6062d6a..d2c917a5d 100644 --- a/util/slicesx/slicesx_test.go +++ b/util/slicesx/slicesx_test.go @@ -44,6 +44,7 @@ func BenchmarkInterleave(b *testing.B) { ) } } + func TestShuffle(t *testing.T) { var sl []int for i := 0; i < 100; i++ { @@ -64,3 +65,23 @@ func TestShuffle(t *testing.T) { t.Errorf("expected shuffle after 10 tries") } } + +func TestPartition(t *testing.T) { + var sl []int + for i := 1; i <= 10; i++ { + sl = append(sl, i) + } + + evens, odds := Partition(sl, func(elem int) bool { + return elem%2 == 0 + }) + + wantEvens := []int{2, 4, 6, 8, 10} + wantOdds := []int{1, 3, 5, 7, 9} + if !reflect.DeepEqual(evens, wantEvens) { + t.Errorf("evens: got %v, want %v", evens, wantEvens) + } + if !reflect.DeepEqual(odds, wantOdds) { + t.Errorf("odds: got %v, want %v", odds, wantOdds) + } +}