vicaba
vicaba

Reputation: 2876

Traverse tree in a functional way

I've implemented a basic mutable tree in Scala and I want to traverse it in a functional way in order to search for an element, but I don't know how to implement it. I want also the algorithm to be tail recursive if possible.

The tree is a structure with a value and a list of leafs that are also trees.

Any help would be appreciated.

This is the code that I have (focus on getOpt method):

package main

import scala.collection.mutable.ListBuffer

sealed trait Tree[Node] {

  val node: Node

  val parent: Option[Tree[Node]]

  val children: ListBuffer[Tree[Node]]

  def add(n: Node): Tree[Node]

  def size: Int

  def getOpt(p: (Node) => Boolean): Option[Tree[Node]]

  override def toString = {
    s"""[$node${if (children.isEmpty) "" else s", Children: $children"}]"""
  }
}

case class ConcreteTree(override val node: Int) extends Tree[Int] {

  override val children = ListBuffer[Tree[Int]]()

  override val parent: Option[Tree[Int]] = None

  override def add(n: Int): ConcreteTree = {
    val newNode = new ConcreteTree(n) {override val parent: Option[Tree[Int]] = Some(this)}
    children += newNode
    newNode
  }

  override def size: Int = {
    def _size(t: Tree[Int]): Int = {
      1 + t.children.foldLeft(0)((sum, tree) => sum + _size(tree))
    }
    _size(this)
  }

  // This method is not correct
  override def getOpt(p: (Int) => Boolean): Option[Tree[Int]] = {
    def _getOpt(tree: Tree[Int], p: (Int) => Boolean): Option[Tree[Int]] = {
      tree.children.map {
        t =>
          if(p(t.node)) Some(t) else t.children.map(_getOpt(_, p))
      }
    }
  }
}

object Main {

  def main(args: Array[String]) {
    val tree1 = ConcreteTree(1)
    val tree2 = tree1.add(2)
    val tree3 = tree1.add(3)

    println(tree1.getOpt(_ == 2))
  }
}

@doertsch answer is the best approach that I have at the moment.

Upvotes: 2

Views: 3110

Answers (4)

Karl Bielefeldt
Karl Bielefeldt

Reputation: 49008

I would actually go for something more flexible and implement a generic function to produce a lazy stream of your flattened tree, then a lot of your later work will become much easier. Something like this:

def traverse[Node](tree: Tree[Node]): Stream[Tree[Node]] = 
  tree #:: (tree.children map traverse).fold(Stream.Empty)(_ ++ _)

Then your getOpt reduces to this:

override def getOpt(p: (Int) => Boolean): Option[Tree[Int]] =
  traverse(tree) find {x => p(x.node)}

Simplifying even further, if you're just interested in the data without the Tree structure, you can just get a stream of nodes, giving you:

def nodes[Node](tree: Tree[Node]): Stream[Node] = traverse(tree) map (_.node)

def getNode(p: (Int) => Boolean): Option[Int] = nodes(tree) find p

This opens other possibilities for very concise and readable code like nodes(tree) filter (_ > 3), nodes(tree).sum, nodes(tree) contains 12, and similar expressions.

Upvotes: 7

Dima
Dima

Reputation: 40500

I think, you are looking for something like this:

@tailrec
def findInTree[T](value: T, stack: List[Node[T]]): Option[Node[T]] = stack match {
   case Nil => None
   case Node(value) :: _ => stack.headOption
   case head :: tail => findInTree(value, head.children ++ tail)
}

This does not traverse parts of the tree twice, and is also tail-recursive. I changed your class definitions a little bit for clarity, but it's the same idea. You call it something like this: findInTree(valueToFind, List(root))

Upvotes: 2

lex82
lex82

Reputation: 11297

When turning a collection into a view, all transformers like map are implemented lazily. This means elements are only processed as required. This should solve your problem:

override def getOpt(p: (Int) => Boolean): Option[Tree[Int]] = {
  if (p(node)) Some(this)
  else children.view.map(_.getOpt(p)).find(_.isDefined).getOrElse(None)
}

So we are mapping (lazily) over the children, turning them into Options of the searched node. Subsequently we find the first such Option being not None. The final getOrElse(None) is required to "flatten" the nested Options since find returns an Option itself.

I didn't actually run the code, so please excuse minor mistakes. However, the general approach should have become clear.

Upvotes: 1

doertsch
doertsch

Reputation: 121

Using the methods exists and find (which every List provides), you can achieve the "finish when a result is found" behaviour. (Although it might be argued that internally, these are not implemented in a totally functional way: https://github.com/scala/scala/blob/5adc400f5ece336f3f5ff19691204975d41e652e/src/library/scala/collection/LinearSeqOptimized.scala#L88)

Your code might look as follows:

case class Tree(nodeValue: Long, children: List[Tree]) {

  def containsValue(search: Long): Boolean =
    search == nodeValue || children.exists(_.containsValue(search))

  def findSubTreeWithNodeValue(search: Long): Option[Tree] =
    if (search == nodeValue)
      Some(this)
    else
      children.find(_.containsValue(search)).
        flatMap(_.findSubTreeWithNodeValue(search))
}

On the last two lines, the find application will return the correct subtree of the current node, in case one exists, and the flatMap part will extract the correct subtree from it recursively, or leave the None result unchanged if the value was not found.

This one, however, has the unlovely characteristic of doing parts of the traversal twice, once for determining whether a result exists, and once for extracting it from the tree that contains it. There might be a way to fix this in a more efficient way, but I can't figure it out at the moment ...

Upvotes: 2

Related Questions