diff --git a/groupby.go b/groupby.go new file mode 100644 index 0000000..e4a554b --- /dev/null +++ b/groupby.go @@ -0,0 +1,15 @@ +package underscore + +// GroupBy splits a slice into a map[K][]V grouped by the result of the iterator function. +func GroupBy[K comparable, V any](values []V, f func(V) K) map[K][]V { + res := make(map[K][]V, 0) + for _, v := range values { + k := f(v) + if r, ok := res[k]; ok { + res[k] = append(r, v) + } else { + res[k] = []V{v} + } + } + return res +} diff --git a/groupby_test.go b/groupby_test.go new file mode 100644 index 0000000..ef4d4e0 --- /dev/null +++ b/groupby_test.go @@ -0,0 +1,22 @@ +package underscore_test + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" + + u "github.com/rjNemo/underscore" +) + +func TestGroupBy(t *testing.T) { + nums := []float64{1.3, 2.1, 2.4} + want := map[int][]float64{ + 1: {1.3}, + 2: {2.1, 2.4}, + } + f := func(n float64) int { + return int(math.Floor(n)) + } + assert.Equal(t, want, u.GroupBy(nums, f)) +}