Efficient way of flattening a recursive data structure in golang

I have a recursive data structure that can contain a few different type of data:

type Data interface{
 // Some methods
}
type Pair struct { // implements Data
  fst Data
  snd Data
}
type Number float64 // implements Data

Now I want to flatten a chain of Pairs into a []Data. However, the Data in the fst field should not be flattened, only data in snd should be flattened. E.g:

chain := Pair{Number(1.0), Pair{Number(2.0), Pair{Number(3.0), nil}}}
chain2 := Pair{Pair{Number(1.0), Number(4.0)}, Pair{Number(2.0), Pair{Number(3.0), nil}}}

becomes:

data := []Data{Number(1.0), Number(2.0), Number(3.0)}
data2 := []Data{Pair{Number(1.0), Number(4.0)}, Number(2.0), Number(3.0)}

My naive approach would be:

var data []Data
chain := Pair{Number(1.0), Pair{Number(2.0), Pair{Number(3.0), nil}}}

for chain != nil {
    data = append(data, chain.fst)
    chain = chain.snd
}

Is there a more efficient approach that can flatten a data structure like the one in the variable chain into an []Data array?

You can use a recursive function. On the way down, add up the number of pairs, at the bottom, allocate the array, and on the way back up, fill the array from back to front.

If you need to support arbitrary trees, you can add a size method to Data, and then do another tree traversal to actually fill the array.

Huh, your naive approach doesn't work for Pairs nested inside fst. If you had chain := Pair{Pair{Number(1.0), Number(2.0)}, Number{3.0}}, it would end up as []Data{Pair{Number(1.0), Number(2.0)}, Number{3.0}}. This is an inherently recursive problem, so why not implement it as such?

I suggest adding a flatten() method to your interface. Pairs can just recursively nest themselves, and Numbers just return their value.

Here's a fully working example with some minimal testing:

package main

import "fmt"

type Data interface {
	flatten() []Data
}

type Pair struct {
	fst Data
	snd Data
}

type Number float64

func (p Pair) flatten() []Data {
	res := []Data{}
	if p.fst != nil {
		res = append(res, p.fst.flatten()...)
	}
	if p.snd != nil {
		res = append(res, p.snd.flatten()...)
	}
	return res
}

func (n Number) flatten() []Data {
	return []Data{n}
}

func main() {
	tests := []Data{
		Pair{Number(1.0), Pair{Number(2.0), Pair{Number(3.0), nil}}},
		Pair{Pair{Number(1.0), Number(2.0)}, Number(3.0)},
		Pair{Pair{Pair{Number(1.0), Number(2.0)}, Pair{Number(3.0), Number(4.0)}}, Pair{Pair{Number(5.0), Number(6.0)}, Number(7.0)}},
		Number(1.0),
	}
	for _, t := range tests {
		fmt.Printf("Original: %v\n", t)
		fmt.Printf("Flattened: %v\n", t.flatten())
	}
}

(This assumes that the top-level input Data is never nil).

The code prints:

Original: {1 {2 {3 <nil>}}}
Flattened: [1 2 3]
Original: {{1 2} 3}
Flattened: [1 2 3]
Original: {{{1 2} {3 4}} {{5 6} 7}}
Flattened: [1 2 3 4 5 6 7]
Original: 1
Flattened: [1]

As suggested, writing a recursive function fits best for this problem. But it's also possible to write a non-recursive version (IMHO recursive version would be more clear):

func flatten(d Data) []Data {
	var res []Data

	stack := []Data{d}

	for {
		if len(stack) == 0 {
			break
		}
		switch x := stack[len(stack)-1].(type) {
		case Pair:
			stack[len(stack)-1] = x.snd
			stack = append(stack, x.fst)
		case Number:
			res = append(res, x)
			stack = stack[:len(stack)-1]
		default:
			if x == nil {
				stack = stack[:len(stack)-1]
			} else {
				panic("INVALID TYPE")
			}
		}
	}

	return res
}