Add AllRepeat function (#16)

Author: alexykot

combinations.go

@@ -27,6 +27,34 @@ func All[T any](set []T) (subsets [][]T) {
 	return subsets
 }
 
+// AllRepeat returns all combinations with repetitions for a given slice,
+// from 1 up to a maximum combination length of m.
+func AllRepeat[T any](set []T, m int) (subsets [][]T) {
+	if m < 1 {
+		return nil
+	}
+
+	var generateCombos func([]T, int)
+	generateCombos = func(current []T, depth int) {
+		if depth == 0 {
+			subset := make([]T, len(current))
+			copy(subset, current)
+			subsets = append(subsets, subset)
+			return
+		}
+
+		for _, item := range set {
+			generateCombos(append(current, item), depth-1)
+		}
+	}
+
+	for length := 1; length <= m; length++ {
+		generateCombos([]T{}, length)
+	}
+
+	return subsets
+}
+
 // Combinations returns combinations of n elements for a given generic array.
 // For n < 1, it equals to All and returns all combinations.
 func Combinations[T any](set []T, n int) (subsets [][]T) {

combinations_test.go

@@ -264,3 +264,218 @@ func TestStringCombinationsN(t *testing.T) {
 		})
 	}
 }
+
+func TestAllWithRepetitions(t *testing.T) {
+	tt := []struct {
+		name string
+		in   []string
+		m    int
+		out  [][]string
+	}{
+		{
+			name: "Empty slice",
+			in:   []string{},
+			m:    1,
+			out:  nil,
+		},
+		{
+			name: "Single item, m = 1",
+			in:   []string{"A"},
+			m:    1,
+			out: [][]string{
+				{"A"},
+			},
+		},
+		{
+			name: "Single item, m = 2",
+			in:   []string{"A"},
+			m:    2,
+			out: [][]string{
+				{"A"},
+				{"A", "A"},
+			},
+		},
+		{
+			name: "Two items, m = 1",
+			in:   []string{"A", "B"},
+			m:    1,
+			out: [][]string{
+				{"A"},
+				{"B"},
+			},
+		},
+		{
+			name: "Two items, m = 2",
+			in:   []string{"A", "B"},
+			m:    2,
+			out: [][]string{
+				{"A"},
+				{"B"},
+				{"A", "A"},
+				{"A", "B"},
+				{"B", "A"},
+				{"B", "B"},
+			},
+		},
+		{
+			name: "Three items, m = 2",
+			in:   []string{"A", "B", "C"},
+			m:    2,
+			out: [][]string{
+				{"A"},
+				{"B"},
+				{"C"},
+				{"A", "A"},
+				{"A", "B"},
+				{"A", "C"},
+				{"B", "A"},
+				{"B", "B"},
+				{"B", "C"},
+				{"C", "A"},
+				{"C", "B"},
+				{"C", "C"},
+			},
+		},
+		{
+			name: "Two items, m = 3",
+			in:   []string{"A", "B"},
+			m:    3,
+			out: [][]string{
+				{"A"},
+				{"B"},
+				{"A", "A"},
+				{"A", "B"},
+				{"B", "A"},
+				{"B", "B"},
+				{"A", "A", "A"},
+				{"A", "A", "B"},
+				{"A", "B", "A"},
+				{"A", "B", "B"},
+				{"B", "A", "A"},
+				{"B", "A", "B"},
+				{"B", "B", "A"},
+				{"B", "B", "B"},
+			},
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			out := AllRepeat(tc.in, tc.m)
+			if !reflect.DeepEqual(out, tc.out) {
+				t.Errorf("error: \nreturn:\t%v\nwant:\t%v", out, tc.out)
+			}
+		})
+	}
+}
+
+func TestAllWithRepetitionsInt(t *testing.T) {
+	tt := []struct {
+		name string
+		in   []int
+		m    int
+		out  [][]int
+	}{
+		{
+			name: "Two items, m = 0",
+			in:   []int{1, 2},
+			m:    0,
+			out:  nil,
+		},
+		{
+			name: "Empty slice",
+			in:   []int{},
+			m:    1,
+			out:  nil,
+		},
+		{
+			name: "Single item, m = 1",
+			in:   []int{1},
+			m:    1,
+			out: [][]int{
+				{1},
+			},
+		},
+		{
+			name: "Single item, m = 2",
+			in:   []int{1},
+			m:    2,
+			out: [][]int{
+				{1},
+				{1, 1},
+			},
+		},
+		{
+			name: "Two items, m = 2",
+			in:   []int{1, 2},
+			m:    2,
+			out: [][]int{
+				{1},
+				{2},
+				{1, 1},
+				{1, 2},
+				{2, 1},
+				{2, 2},
+			},
+		},
+		{
+			name: "Three items, m = 1",
+			in:   []int{1, 2, 3},
+			m:    1,
+			out: [][]int{
+				{1},
+				{2},
+				{3},
+			},
+		},
+		{
+			name: "Three items, m = 2",
+			in:   []int{1, 2, 3},
+			m:    2,
+			out: [][]int{
+				{1},
+				{2},
+				{3},
+				{1, 1},
+				{1, 2},
+				{1, 3},
+				{2, 1},
+				{2, 2},
+				{2, 3},
+				{3, 1},
+				{3, 2},
+				{3, 3},
+			},
+		},
+		{
+			name: "Two items, m = 3",
+			in:   []int{1, 2},
+			m:    3,
+			out: [][]int{
+				{1},
+				{2},
+				{1, 1},
+				{1, 2},
+				{2, 1},
+				{2, 2},
+				{1, 1, 1},
+				{1, 1, 2},
+				{1, 2, 1},
+				{1, 2, 2},
+				{2, 1, 1},
+				{2, 1, 2},
+				{2, 2, 1},
+				{2, 2, 2},
+			},
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			out := AllRepeat(tc.in, tc.m)
+			if !reflect.DeepEqual(out, tc.out) {
+				t.Errorf("error: \nreturn:\t%v\nwant:\t%v", out, tc.out)
+			}
+		})
+	}
+}