Go Tree Flattener

main.go

package main

import (
	"fmt"
)

// Nums is a wrapper around []int
type Nums []int

// Nodes represents an arbitrarily nested tree structure of Node structsß
type Nodes []Node

// Node represent a leaf-level item in the tree
type Node struct {
	Val      Nums
	Children Nodes
}

func main() {
	nodes := Nodes{}

	n1 := Node{Nums{3}, nil}
	n2 := Node{Nums{1, 2, 4}, Nodes{n1}}

	nodes = append(nodes, n2)

	fmt.Println(FlattenNodes(nodes))
}

// FlattenNodes will recursively flatten Nodes{} int Nums{}.
func FlattenNodes(nodes Nodes) Nums {
	nums := Nums{}

	for _, node := range nodes {		
		for _, num := range node.Val {
			nums = append(nums, num)
		}
		
		if node.Children != nil {
			nums = append(nums, FlattenNodes(node.Children)...)
		}
	}

	return nums
}

main_test.go

package main

import (
	"reflect"
	"testing"
)

func TestFlattenNodes(t *testing.T) {
	type args struct {
		nodes Nodes
	}
	tests := []struct {
		name string
		args args
		want Nums
	}{
		{
			"single level, single item",
			args{Nodes{Node{[]int{1, 2, 3}, nil}}},
			Nums{1, 2, 3},
		},
		{
			"single level, multiple items",
			args{Nodes{Node{[]int{1, 2, 3}, nil}, Node{[]int{4, 5, 6}, nil}}},
			Nums{1, 2, 3, 4, 5, 6},
		},
		{
			"nested, single item",
			args{Nodes{Node{[]int{1, 2, 3}, Nodes{Node{[]int{4, 5, 6}, nil}}}}},
			Nums{1, 2, 3, 4, 5, 6},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := FlattenNodes(tt.args.nodes); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("FlattenNodes() = %v, want %v", got, tt.want)
			}
		})
	}
}