Combination Algorithm in Golang - Combinations from n arrays picking one element from each array

main_test.go

//source: https://www.geeksforgeeks.org/combinations-from-n-arrays-picking-one-element-from-each-array/

package main

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func Combination(arr [][]int) [][]int {
	var combination [][]int
	n := len(arr)

	indices := make([]int, n)

	for {
		var innerArray []int

		for i := 0; i < n; i++ {
			item := arr[i][indices[i]]
			innerArray = append(innerArray, item)
		}

		combination = append(combination, innerArray)

		next := n - 1

		for next >= 0 && (indices[next]+1 >= len(arr[next])) {
			next--
		}

		if next < 0 {
			break
		}

		indices[next]++

		for i := next + 1; i < n; i++ {
			indices[i] = 0
		}
	}

	return combination

}

func TestCombination(t *testing.T) {

	t.Run("2 arrays", func(t *testing.T) {
		arr := [][]int{
			{1, 2},
			{3, 4},
		}

		result := Combination(arr)

		expectedResult := [][]int{
			{1, 3},
			{1, 4},
			{2, 3},
			{2, 4},
		}

		assert.Len(t, result, 4)

		for i, innerArray := range result {
			for j, ia := range innerArray {
				assert.Equal(t, expectedResult[i][j], ia)
			}
		}
	})

	t.Run("3 arrays", func(t *testing.T) {
		arr := [][]int{
			{1},
			{2, 3, 4},
			{5},
		}

		result := Combination(arr)

		expectedResult := [][]int{
			{1, 2, 5},
			{1, 3, 5},
			{1, 4, 5},
		}

		assert.Len(t, result, 3)

		for i, innerArray := range result {
			for j, ia := range innerArray {
				assert.Equal(t, expectedResult[i][j], ia)
			}
		}
	})
}