CODEWITHSUNDEEP
scala
PepeHands
PepeHands

Reputation: 1406

Scala: use for comprehension to iterate over tree node parents

Suppose that I have some Tree class which internally uses Nodes:

class Tree {
    val root = Node(null)
    case class Node(parent: Node) {
        val children: List[Node]
    }
}

now suppose that I want to get all ancestors of some node applying some filter function. The obvious way to do it is something like this:

def getAncsWithFilter(filter: (Node) => Boolean): List[Node] {
    var curNode = this.parent
    var res: List[Node] = Nil
    while (curNode != null) {
        if (filter(curNode)) {
            res = res :: curNode
        }
        curNode = curNode.parent
    }
    res
}

What I don't like in this solution is that I use vars and this code seems too imperative and ugly. What I want to be able to write is something like this

def getAncsWithFilter(filter: (Node) => Boolean): List[Node] {
    val curNode = this
    for (curNode <- curNode.parent if filter(curNode))
        yield curNode
}

Is it possible is Scala? Currently, if I write code like this, it will generate some error complaining about absence of withFilter method. If I extend my Node class with Traversable[Node] then code above will return Traversable instead of List.

Upvotes: 2

Views: 522

Answers (2)

adrice727
adrice727

Reputation: 1492

A for-comprehension would require an implementation of flatMap on Node. Another option is to use recursion:

def getAncsWithFilter(filter: (Node) => Boolean, ancs: List[Node] = List()): List[Node] = {
  (parent, filter(parent)) match {
    case (null, _)  => ancs
    case (_, true)  => parent.getAncsWithFilter(filter, parent :: ancs)
    case (_, false) => parent.getAncsWithFilter(filter, ancs)
  }
}

Upvotes: 0

Kolmar
Kolmar

Reputation: 14224

The problem here is that curNode.parent returns a single Node, not a sequence of all ancestors, so you are not iterating over anything.

You can implement a separate method ancestors, that would return a sequence of all ancestors and you would be able to iterate through it.

A quick way to get something iterable with ancestors, is to use Iterator.iterate:

def ancestors: Iterator[Node] = 
  Iterator.iterate(this.parent)(_.parent).takeWhile(_ != null)

The full implementation of getAncsWithFilter will then become:

def getAncsWithFilter(filter: (Node) => Boolean): List[Node] =
  ancestors.filter(filter).toList

You don't even have to use a for-comprehension (you can still use it of course, but in that case the code will become more complex in my opinion).

Upvotes: 4

Related Questions