flatten.go

package flatten_test

import (
	"github.com/stretchr/testify/require"
	"math/rand"
	"testing"
)

// Golang has no multi-dimensional arrays, so I used a tree
type Node struct {
	Nodes []Node // Pointer to a next array if not nil
	Value int    // Integer value
}

// Flatten tree to a slice of integers
func Flatten(n []Node) []int {
	ret := make([]int, 0, len(n))
	for _, v := range n {
		if v.Nodes == nil {
			ret = append(ret, v.Value)
		} else {
			ret = append(ret, Flatten(v.Nodes)...)
		}
	}

	return ret
}

// pre-generation of a random tree for benchmarks
var bgen []Node

func init() {
	// generating test tree
	bgen = gen(100, 0)
}

func gen(n int, l int) []Node {
	ret := make([]Node, 0, n)
	for i := 0; i < n; i++ {
		if rand.Intn(100) <= 90 {
			ret = append(ret, Node{Value: rand.Intn(100000)})
		} else if l < 3 {
			ret = append(ret, gen(n, l+1)...)
		}
	}

	return ret
}

func TestFlatten(t *testing.T) {
	s := []Node{
		{Value: 1},
		{Value: 2},
		{
			Nodes: []Node{
				{Value: 3},
				{Nodes: []Node{
					{Value: 4},
				}},
			},
		},
		{Value: 5},
	}

	require.Equal(t, []int{1, 2, 3, 4, 5}, Flatten(s))
}

func BenchmarkFlatten(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Flatten(bgen)
	}
}