scala count number of nodes in a tree - scala

Question

I have a tree defined as
sealed trait Tree[+A]
case class Leaf[A](value: A) extends Tree[A]
case class Branch[A](left: Tree[A], right: Tree[A]) extends Tree[A]
I would like a function to count the number of nodes in the tree and a function that counts the number of leafs
(I wrote a function to counts the number of leafs but I am not satisfied of the method that I am using. I would like a more "functional" implementation of that one).
val t1 = Branch(
Branch(
Leaf(12),
Branch(
Leaf(3),
Leaf(4))),
Leaf(8)) //> t1 : trees.Branch[Int] = Branch(Branch(Leaf(12),Branch(Leaf(3),Leaf(4))),Le
//| af(8))
def n_nodes(t:Tree[Int]):Int = {
var s = 0
def n_nodesAcc(t:Tree[Int]):Unit = {
t match {
case Branch(left, right) =>{
n_nodesAcc(left)
n_nodesAcc(right )
}
case Leaf(v) => {s = s+ 1}
}
}
n_nodesAcc(t)
s
} //> n_nodes: (t: trees.Tree[Int])Int
n_nodes(t1) //> res0: Int = 4
(this is an exercise)

Answer 1

You might write a recursive method to count the leaves:
def countLeaves[A](tree: Tree[A]): Int = tree match {
case l:Leaf[A] => 1
case b:Branch[A] => countLeaves(b.left) + countLeaves(b.right)
}
Or count all the nodes:
def countNodes[A](tree: Tree[A]): Int = tree match {
case l:Leaf[A] => 1
case b:Branch[A] => 1 + countLeaves(b.left) + countLeaves(b.right)
}
You could also write a similar method to just get all the leaves, and you could have more flexibility to do different functions later:
def getLeaves[A](tree: Tree[A]): Seq[Leaf[A]] = tree match {
case l:Leaf[A] => Seq(l)
case b:Branch[A] => getLeaves(b.left) ++ getLeaves(b.right)
}
And then count with getLeaves(t1).length. Or, similarly, get all the nodes:
def getNodes[A](tree: Tree[A]): Seq[Tree[A]] = tree match {
case l:Leaf[A] => Seq(l)
case b:Branch[A] => Seq(b) ++ getNodes(b.left) ++ getNodes(b.right)
}