perf: add precision parameter to xslices.Average()

fufuok · Jun 3, 2024 · eee33db · eee33db
1 parent e8b387e
commit eee33db
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/generic/xslices/utils.go b/generic/xslices/utils.go
@@ -27,8 +27,8 @@ func ToString[T any](xs []T, sep string) string {
 	return strings.Join(ss, sep)
 }
 
-// Average 求数字切片的平均值
-func Average[T generic.Numeric](xs []T) float64 {
+// Average 求数字切片的平均值, 可选指定保留小数位数
+func Average[T generic.Numeric](xs []T, precision ...int) float64 {
 	switch len(xs) {
 	case 0:
 		return 0
@@ -41,5 +41,9 @@ func Average[T generic.Numeric](xs []T) float64 {
 	for _, x := range xs {
 		m += x
 	}
-	return float64(m) / float64(len(xs))
+	v := float64(m) / float64(len(xs))
+	if len(precision) > 0 && precision[0] > 0 {
+		return utils.Round(v, precision[0])
+	}
+	return v
 }
diff --git a/generic/xslices/utils_test.go b/generic/xslices/utils_test.go
@@ -24,4 +24,9 @@ func TestAverage(t *testing.T) {
 	assert.Equal(t, 2.5, Average(xs))
 	assert.Equal(t, float64(0), Average([]float32{0.0}))
 	assert.Equal(t, 0.6, Average([]float64{0.0, 1.2}))
+
+	fs := []float64{0.0, 0.5, 0.5}
+	assert.Equal(t, 0.33, Average(fs, 2))
+	fs = []float64{0.0, 1, 1}
+	assert.Equal(t, 0.667, Average(fs, 3))
 }