Scala - flattening a tree structure

Question

I have a tree structure I'm receiving from a Java library. I am trying to flatten it since I'm interested only in the "key" values of the tree. The tree is made up of zero or more of the following classes:

class R(val key: String, val nodes: java.util.List[R]) {}

with an empty nodes list representing the end of a branch. A sample can be build via this code:

val sample =  List[R](
  new R("1",  List[R](
    new R("2",  List[R]().asJava),
    new R("3",  List[R](new R("4",  List[R]().asJava))
      .asJava)).asJava)).asJava

I am having trouble writing both a correct method, and an efficient method. This is what I have so far:

def flattenTree(tree: List[R]): List[String] = {
  tree.foldLeft(List[String]())((acc, x) => 
             x.key :: flattenTree(x.nodes.asScala.toList))
}

However when I run this code, as inefficient as it may be, I still get it incorrect. My result ends up being:

>>> flattenTree(sample.asScala.toList)
res0: List[String] = List(1, 3, 4)

which means for some reason I lost the node with key "2".

Can someone recommend a correct and more efficient way of flattening this tree?

Answer 1

You can define a function to flatten an R object with flatMap:

// required to be able to use flatMap on java.util.List
import scala.collection.JavaConversions._

def flatten(r: R): Seq[String] = {
  r.key +: r.nodes.flatMap(flatten)
}

And a function to flatten a sequence of those:

def flattenSeq(l: Seq[R]): Seq[String] = l flatMap flatten

r.nodes.flatMap(flatten) is a Buffer, so prepending to it is not efficient. It becomes quadratic complexity. So, if the order is not important is more efficient to append: def flatten(r: R): Seq[String] = r.nodes.flatMap(flatten) :+ r.key

Answer 2

You are failing to add in the accumulated keys on each successive call. Try the following:

def flattenTree(tree: List[R]): List[String] = {
  tree.foldLeft(List[String]())((acc, x) =>
             x.key :: flattenTree(x.nodes.asScala.toList) ++ acc)
}

which generates the result: List(1, 3, 4, 2),

or, if proper ordering is important:

def flattenTree(tree: List[R]): List[String] = {
  tree.foldLeft(List[String]())((acc, x) =>
             acc ++ (x.key :: flattenTree(x.nodes.asScala.toList)))
}

which generates the result: List(1, 2, 3, 4)