diff --git a/flatmap.go b/flatmap.go index 1ce9c80..352176d 100644 --- a/flatmap.go +++ b/flatmap.go @@ -2,7 +2,8 @@ package underscore // Flatmap flatten the input slice element into the new slice. FlatMap maps every element with the help of a mapper function, then flattens the input slice element into the new slice. func Flatmap[T any](values []T, mapper func(n T) []T) []T { - res := make([]T, 0) + // Estimate capacity: assume average of 2-3 items per element + res := make([]T, 0, len(values)*2) for _, v := range values { vs := mapper(v) res = append(res, vs...) diff --git a/flatmap_test.go b/flatmap_test.go index ff5ca41..714c0b3 100644 --- a/flatmap_test.go +++ b/flatmap_test.go @@ -15,3 +15,16 @@ func TestFlatmap(t *testing.T) { assert.Equal(t, want, u.Flatmap(nums, transform)) } + +func BenchmarkFlatmap(b *testing.B) { + data := make([]int, 100) + for i := range data { + data[i] = i + } + mapper := func(n int) []int { return []int{n, n * 2, n * 3} } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + u.Flatmap(data, mapper) + } +}