diff --git a/common/graph/choose.go b/common/graph/choose.go index 08a6ce0dc49..d7df824868f 100644 --- a/common/graph/choose.go +++ b/common/graph/choose.go @@ -59,7 +59,14 @@ func choose(n int, targetAmount int, i int, currentSubGroup []int, subGroups *or return } // We either pick the current element - choose(n, targetAmount, i+1, append(currentSubGroup, i), subGroups) + choose(n, targetAmount, i+1, concatInts(currentSubGroup, i), subGroups) // Or don't pick it choose(n, targetAmount, i+1, currentSubGroup, subGroups) } + +func concatInts(a []int, elements ...int) []int { + var res []int + res = append(res, a...) + res = append(res, elements...) + return res +} diff --git a/common/graph/choose_test.go b/common/graph/choose_test.go index 442b4f0f178..45c0cc2f70c 100644 --- a/common/graph/choose_test.go +++ b/common/graph/choose_test.go @@ -7,6 +7,7 @@ SPDX-License-Identifier: Apache-2.0 package graph import ( + "fmt" "testing" "github.com/stretchr/testify/assert" @@ -24,3 +25,32 @@ func TestCombinationsExceed(t *testing.T) { // N < K returns false assert.False(t, CombinationsExceed(20, 30, 0)) } + +func TestChooseKoutOfN(t *testing.T) { + expectedSets := indiceSets{ + &indiceSet{[]int{0, 1, 2, 3}}, + &indiceSet{[]int{0, 1, 2, 4}}, + &indiceSet{[]int{0, 1, 2, 5}}, + &indiceSet{[]int{0, 1, 3, 4}}, + &indiceSet{[]int{0, 1, 3, 5}}, + &indiceSet{[]int{0, 1, 4, 5}}, + &indiceSet{[]int{0, 2, 3, 4}}, + &indiceSet{[]int{0, 2, 3, 5}}, + &indiceSet{[]int{0, 2, 4, 5}}, + &indiceSet{[]int{0, 3, 4, 5}}, + &indiceSet{[]int{1, 2, 3, 4}}, + &indiceSet{[]int{1, 2, 3, 5}}, + &indiceSet{[]int{1, 2, 4, 5}}, + &indiceSet{[]int{1, 3, 4, 5}}, + &indiceSet{[]int{2, 3, 4, 5}}, + } + require.Equal(t, indiceSetsToStrings(expectedSets), indiceSetsToStrings(chooseKoutOfN(6, 4))) +} + +func indiceSetsToStrings(sets indiceSets) []string { + var res []string + for _, set := range sets { + res = append(res, fmt.Sprintf("%v", set.indices)) + } + return res +}