FlattenArray

flatten.go

package flatten

// FlattenArray flattens the array input, and places inputues into the array output
// When a nested array is encountered, FlattenArray will recurse into itself
//
// In the case that an element in input is not either int or []int, FlattenArray returns nil
func FlattenArray(input []interface{}) []int {
	var output []int
	for _, val := range input {
		switch i := val.(type) {
		case int:
			output = append(output, i)
		case []interface{}:
			output = append(output, FlattenArray(i)...)
		default:
			return nil
		}
	}
	return output
}

flatten_test.go

package flatten

import (
	"reflect"
	"testing"
)

type TestTable struct {
	ExpectOutput []int
	Data         []interface{}
}

var testData = []TestTable{
	TestTable{
		ExpectOutput: nil,
		Data:         []interface{}{true, 2, []interface{}{[]interface{}{3}}, "a"},
	},
	TestTable{
		ExpectOutput: []int{1, 2, 3, 4, 5},
		Data:         []interface{}{1, 2, []interface{}{[]interface{}{3}}, 4, 5},
	},

	TestTable{
		ExpectOutput: []int{1, 2, 3, 4, 5, 6},
		Data: []interface{}{
			[]interface{}{
				[]interface{}{
					[]interface{}{1},
				},
			},
			[]interface{}{
				[]interface{}{
					[]interface{}{
						[]interface{}{2},
					},
				},
			},
			[]interface{}{
				[]interface{}{
					[]interface{}{3},
				},
			},
			[]interface{}{
				[]interface{}{
					[]interface{}{4},
				},
			},
			[]interface{}{
				[]interface{}{
					[]interface{}{5},
				},
			},
			6,
		},
	},

	TestTable{
		ExpectOutput: []int{1, 2, 3, 4},
		Data:         []interface{}{1, 2, 3, 4},
	},
}

func BenchmarkFlattenArray(b *testing.B) {
	b.ReportAllocs()
	for _, data := range testData {
		for i := 0; i < b.N; i++ {
			_ = FlattenArray(data.Data)
		}
	}
}

func TestFlattenArray(t *testing.T) {
	var output []int
	for i, data := range testData {
		output = FlattenArray(data.Data)
		if reflect.DeepEqual(output, data.ExpectOutput) == false {
			t.Fatalf("Test %d failed: Expected output: %v got: %v\n", i, data.ExpectOutput, output)
		}
	}
}