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 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)

Answer 2

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 3

Convert each R to a Scalaz Tree , and call flatten to do a pre-order traversal.

import scala.collection.JavaConversions._
import scalaz._

def rTree(r: R): Tree[String] =
  Tree.node(r.key, r.nodes.toStream.map(rTree))

sample.flatMap(r => rTree(r).flatten): Seq[String]
// List(1, 2, 3, 4)

Edit: Unfortunately, due to a bug in scalaz as of version 7.1.1, this causes a stack overflow for wide trees.

Answer 4

What about using Stream s like scalaz does:

def flatten(rootElem: R): Stream[String] = {
  def flatten0(elem: R, xs: Stream[String]): Stream[String] =
    Stream.cons(elem.key, elem.nodes.foldLeft(xs)((acc, x) => flatten0(x, acc)))

  flatten0(rootElem, Stream.empty)
}