diff --git a/combinations.go b/combinations.go index a17d40e..7499712 100644 --- a/combinations.go +++ b/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) { diff --git a/combinations_test.go b/combinations_test.go index 9cae746..7695aa7 100644 --- a/combinations_test.go +++ b/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) + } + }) + } +}