CODEWITHSUNDEEP
scala
Bernie Wong
Bernie Wong

Reputation: 701

How to traverse an n-ary tree as an iterator or stream

I would like to create a function in my Tree class to traverse an n-ary Tree[T] to get back a tuple with (level, T) so that users of that Tree can do things like 

tree.traverse.foreach{ case(l, t) => printf("%-30s %s", "  "*l + t.id, info(t))

to get a report like this 

rootnode                       
  n1.0                         
    n1.0.1                     >> n1.0.1
    n1.0.2                     >> n1.0.2
    n1.0.3                     >> n1.0.3
  n3.0                         
    n3.0.1                     >> n1.0.1
      n3.0.1.1                 >> n1.0.1

The tree node is 

class TreeNode[T](val id: String, 
          private var _data: Option[T],
          private var _parent: String, 
          private var _treeNodeType: TreeNodeType) {
  private var _children: Set[String] = Set()
  ...
}

I can traverse the tree recursively using the traverse or traversef

class Tree[T] {
  private val ROOT = "rootnode"
  val rootNode = new TreeNode(ROOT, None.asInstanceOf[Option[T]], "", ROOTNODETYPE)
  var hmTree: HashMap[String, TreeNode[T]] = HashMap(ROOT -> rootNode)

def traverse: Unit = {
  def iter(s: String, l: Int): Unit = {
    val n = hmTree(s)
    printf("%-30s\n", "  "*l + n.id)
    n.children.foreach(c => iter(c, l+1))
  }
  iter(ROOT, 0)
}

def traversef(f: Option[T] => String): Unit = {
  def iter(s: String, l: Int): Unit = {
    val n = hmTree(s)
    printf("%-30s %s\n", "  "*l + n.id, f(n.data))
    n.children.foreach(c => iter(c, l+1))
  }
  iter(ROOT, 0)
}

...

I looked at http://www.scala-lang.org/old/node/11275.html (Question: How to implement lazy traversal of n-ary tree using Streams?) but couldn't get the code to work. The thing that tripped me was how to Stream.cons the children.

I am ok with either a stream or an iterator. Key is to define the traverse method in the Tree class and use it outside.

Thanks in advance.

Update

Many thanks to @Archeg -- here is the working traverse

def traverse(t: TreeNode[T], depth: Int = 0): Stream[(TreeNode[T], Int)] = {  
  (t, depth) #:: t.children.foldLeft(Stream.empty[(TreeNode[T], Int)]) { 
          case (aggr, el) => aggr #::: traverse(hmTree(el), depth+ 1)   
  }
}

tree.traverse(tree.rootNode).
  foreach{ case(t, l) => printf("%-30s %s\n", "  "*l + t.id, t.id) }

Upvotes: 3

Views: 742

Answers (1)

Archeg
Archeg

Reputation: 8462

I've made it simple so it would be easy to figure it out:

 case class Tree[T](data: T, children: List[Tree[T]])

 def traverse[T](t: Tree[T]): Stream[T] =
    t.data #:: t.children.foldLeft(Stream.empty[T])((aggr, el) => aggr #::: traverse(el))

Upd

A slightly modified version which provides you with indentation:

def traverseInd[T](t: Tree[T], depth: Int = 0): Stream[(T, Int)] =
    (t.data, depth) #:: t.children.foldLeft(Stream.empty[(T, Int)]) { 
            case (aggr, el) => aggr #::: traverseInd(el, depth+ 1)
    }

Upvotes: 2

Related Questions