[SOLVED] Traversing Scalaz Tree

Traversing Scalaz Tree

I'm trying to understand the scalaz tree structure and am having some difficulty!

First I've defined a tree:

val tree: Tree[Int] =
      1.node(
        2.leaf,
        3.node(
          4.leaf,
          5.leaf))

So far using TreeLoc I've worked out how to find the first element that matches some predicate. E.g. to find the first node where the value is 3:

tree.loc.find(x => x.getLabel == 3)

My next challenge was to try and find all nodes that match some predicate. For example I would like to find all leaf nodes (which should be pretty easy using TreeLoc and isLeaf). Unfortunately I can't for the life of me work out how to walk the tree to do this.

Edit: Sorry I don't think I was clear enough in my original question. To be clear I want to walk the tree in such a way that I have information about the Node available to me. Flatten, foldRight etc just allow me to operate on [Int] whereas I want to be able to operate on Tree[Int] (or TreeLoc[Int]).

Solution

Having a look to how find is implemented in scalaz, my suggestion is to implement something like:

implicit class FilterTreeLoc[A](treeLoc: TreeLoc[A]){
  def filter(p: TreeLoc[A] => Boolean): Stream[TreeLoc[A]] =
    Cobind[TreeLoc].cojoin(treeLoc).tree.flatten.filter(p)
}

It behaves like the find but it gives you back instead a Stream[TreeLoc[A]] instead of an Option[TreeLoc[A]].

You can use this as tree.loc.filter(_.isLeaf) and tree.loc.filter(_.getLabel == 3).

Note: the use of an implicit class can be obviously avoided if you prefer to have this declared as a method instead.