slices: + ToMultimapWithKeys

mikhalytch · Jul 8, 2024 · bd134ad · bd134ad
1 parent 6255fab
commit bd134ad
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/slices/map.go b/slices/map.go
@@ -55,3 +55,19 @@ func ToMapWithKeys[K comparable, V any](keyProducer func(V) K) funcs.Applier[[]V
 		return res
 	}
 }
+
+func ToMultimapWithKeys[K comparable, V any](keyProducer func(V) K) funcs.Applier[[]V, map[K][]V] {
+	return func(vs []V) map[K][]V {
+		res := make(map[K][]V, 0)
+
+		for _, v := range vs {
+			key := keyProducer(v)
+
+			ex := res[key]
+			ex = append(ex, v)
+			res[key] = ex
+		}
+
+		return res
+	}
+}
diff --git a/slices/map_test.go b/slices/map_test.go
@@ -55,3 +55,15 @@ func TestToMapWithKeys(t *testing.T) {
 	require.Equal(t, map[string]int{"a": 0, "b": 1},
 		slices.ToMapWithKeys(func(v int) string { return string([]byte{byte(v + 'a')}) })([]int{0, 1}))
 }
+
+func TestToMultimapWithKeys(t *testing.T) {
+	require.Empty(t, slices.ToMultimapWithKeys(func(_ int) string { return "" })([]int{}))
+	require.Empty(t, slices.ToMultimapWithKeys(func(_ int) string { return "" })(nil))
+
+	require.Equal(t,
+		map[string][]int{"1": {1}},
+		slices.ToMultimapWithKeys(strconv.StoA[int])([]int{1}))
+	require.Equal(t,
+		map[string][]int{"1": {1, 1}, "2": {2, 2, 2}},
+		slices.ToMultimapWithKeys(strconv.StoA[int])([]int{1, 2, 2, 2, 1}))
+}