mapx 增加 merge 方法

Author: Jared-lu

mapx/map.go

@@ -72,3 +72,31 @@ func ToMap[K comparable, V any](keys []K, values []V) (m map[K]V, err error) {
 	}
 	return
 }
+
+// Merge 如果同一个键对应多个值，则会使用最后一个值
+func Merge[K comparable, V any](ms ...map[K]V) map[K]V {
+	fc := func(val1, val2 V) V {
+		return val2
+	}
+	return MergeFunc(fc, ms...)
+}
+
+// MergeFunc 如果同一个键对应多个值，需要指定合并方式
+func MergeFunc[K comparable, V any](mergeFunc func(val1, val2 V) V, ms ...map[K]V) map[K]V {
+	total := 0
+	for _, m := range ms {
+		total += len(m)
+	}
+	merged := make(map[K]V, total)
+
+	for _, m := range ms {
+		for k, v := range m {
+			if value, ok := merged[k]; ok {
+				merged[k] = mergeFunc(value, v)
+			} else {
+				merged[k] = v
+			}
+		}
+	}
+	return merged
+}

mapx/map_test.go

@@ -208,3 +208,54 @@ func TestToMap(t *testing.T) {
 		assert.Equal(t, c.result, result)
 	}
 }
+
+func TestMerge(t *testing.T) {
+	input1 := []map[int]int{
+		{1: 1, 2: 2},
+		{3: 3, 4: 4},
+	}
+	want1 := map[int]int{
+		1: 1, 2: 2, 3: 3, 4: 4,
+	}
+	got := Merge(input1...)
+	assert.Equal(t, want1, got)
+
+	input2 := []map[int]int{
+		{1: 1, 2: 2},
+		{3: 3, 4: 4},
+		{2: 3, 4: 5},
+	}
+	want2 := map[int]int{
+		1: 1, 2: 3, 3: 3, 4: 5,
+	}
+	got = Merge(input2...)
+	assert.Equal(t, want2, want2)
+}
+
+func TestMergeFunc(t *testing.T) {
+	input1 := []map[int]int{
+		{1: 1, 2: 2},
+		{1: 3, 2: 4},
+		{3: 3, 4: 4},
+	}
+	want1 := map[int]int{
+		1: 4, 2: 6, 3: 3, 4: 4,
+	}
+	got := MergeFunc(func(val1, val2 int) int {
+		return val1 + val2
+	}, input1...)
+	assert.Equal(t, want1, got)
+
+	input2 := []map[int]int{
+		{1: 1, 2: 2},
+		{1: 3, 2: 4},
+		{3: 3, 4: 4},
+	}
+	want2 := map[int]int{
+		1: 3, 2: 8, 3: 3, 4: 4,
+	}
+	got = MergeFunc(func(val1, val2 int) int {
+		return val1 * val2
+	}, input2...)
+	assert.Equal(t, want2, got)
+}