GroupBy function for arrays/slices

GroupBy applies a function to each element in the input slice and groups the elements by the function's return value, returning a map from return values to slices of elements. The function f takes an element of type A and returns a key of type K.

func GroupBy[A any, K comparable](input []A, f func(A) K) map[K][]A {
	result := make(map[K][]A)
	for _, v := range input {
		key := f(v)
		result[key] = append(result[key], v)
	}
	return result
}

func TestGroupBy(t *testing.T) {
	input := []string{"apple", "banana", "cherry", "date", "elderberry"}
	f := func(a string) int {
		return len(a)
	}
	expected := map[int][]string{5: {"apple"}, 6: {"banana", "cherry"}, 4: {"date"}, 10: {"elderberry"}}
	result := util.GroupBy(input, f)
	if !reflect.DeepEqual(expected, result) {
		t.Errorf("Expected %v, got %v", expected, result)
	}
}