Folding in binary tree

Question

I implemented a fold method in my scala code. I can use it to determine the size and the depth of the tree. Now I'd like to implement the methods max and map which should work like here. The difference is, that the value is saved in the branch instead of the leaf. Here's my code so far:

sealed trait Tree[+A] {
  def size(): Int = fold(this)(() => 1)((l, r, v) => 1 + l + r)

  def depth(): Int = fold(this)(() => 0)((left: Int, right: Int, v: A) => 1 + (left max right))

  def fold[X, B](t: Tree[X])(f: () => B)(g: (B, B, X) => B): B = t match {
    case Leaf() => f()
    case Branch(left, right, v) => g(fold(left)(f)(g), fold(right)(f)(g), v)
  }
}

case class Leaf[A]() extends Tree[A]

case class Branch[A](left: Tree[A], right: Tree[A], v: A) extends Tree[A]

Can anybody help me with that?

Answer 1

Firstly, you can put fold to companion object of Tree (since fold doesn't use this i.e. is "static").

Secondly, start with implementing map for your case

def map[A,B](t: Tree[A])(f: A => B): Tree[B] = t match {
  case Leaf() => ???_1
  case Branch(l, r, v) => ???_2 // here l, r are Tree[A]'s i.e. subtrees of t
}

looking how it's implemented for trees with values in leaves.

Then replace this implementation with the one using fold

def map[B](f: A => B): Tree[B] = 
  fold[A, Tree[B]](this)(() => ???_1_)((l: Tree[B], r: Tree[B], v: A) => ???_2_))
  // here l, r are Tree[B]'s i.e. results of map for subtrees

???_1_ and ???_2_ are ???_1, ???_2, where you replace recursive call of map with l, r, which are results of the recursive call. So ???_1_ is exactly ???_1.