Binary Tree in Scala

tree.sh

#!/bin/sh
exec scala -savecompiled "$0" "$@"

# Output
# --------------------------------------------------------
# Node(1,Node(2,Leaf(4),Leaf(5)),Node(3,Leaf(6),Leaf(7)))
# ========================================================
# size        7
# height      3
# leafCount   4
# leaves      List(4, 5, 6, 7)
# inOrder     List(4, 2, 5, 1, 6, 3, 7)
# preOrder    List(1, 2, 4, 5, 3, 6, 7)
# postOrder   List(4, 5, 2, 6, 7, 3, 1)
# levelOrder  List(1, 2, 3, 4, 5, 6, 7)
# --------------------------------------------------------
# Node('F,Node('B,Leaf('A),Node('D,Leaf('C),Leaf('E))),Node('G,Empty,Node('I,Leaf('H),Empty))) 
# ========================================================
# size        9
# height      4
# leafCount   4
# leaves      List('A, 'C, 'E, 'H)
# inOrder     List('A, 'B, 'C, 'D, 'E, 'F, 'G, 'H, 'I)
# preOrder    List('F, 'B, 'A, 'D, 'C, 'E, 'G, 'I, 'H)
# postOrder   List('A, 'C, 'E, 'D, 'B, 'H, 'I, 'G, 'F)
# levelOrder  List('F, 'B, 'G, 'A, 'D, 'I, 'C, 'E, 'H)
!#

import annotation.tailrec

trait Tree[+A] {

  private def fold[A, B](t: Tree[A], z: B)(f: (B, A) => B)(g: (B, B, B) => B): B = t match {
    case Empty => z
    case Leaf(v) => f(z, v)
    case Node(v, l, r) => g(f(z, v), fold(l, z)(f)(g), fold(r, z)(f)(g))
  }

  def size: Int           = fold(this, 0)((z, v) => 1)((_, l, r) => l + r + 1)
  def leafCount: Int      = fold(this, 0)((z, v) => 1)((_, l, r) => l + r)
  def height: Int         = fold(this, 0)((z, v) => 1)((_, l, r) => (l max r) + 1)

  def leaves: List[A]     = fold(this, List[A]())((z, v) => v :: z)((v, l, r) => l ::: r)
  def inOrder: List[A]    = fold(this, List[A]())((z, v) => v :: z)((v, l, r) => l ::: v ::: r)
  def preOrder: List[A]   = fold(this, List[A]())((z, v) => v :: z)((v, l, r) => v ::: l ::: r)
  def postOrder: List[A]  = fold(this, List[A]())((z, v) => v :: z)((v, l, r) => l ::: r ::: v)

  def levelOrder: List[A] = {
    def next(tree: Tree[A]): List[Tree[A]] = tree match {
      case Node(_, l, r) => List(l, r)
      case _ => Nil
    }

    def flatten(tree: List[Tree[A]]): List[A] = tree flatMap {
      case Empty => None
      case Leaf(v) => Some(v)
      case Node(v, _, _) => Some(v)
    }

    @tailrec
    def loop(tree: List[Tree[A]], acc: List[A]): List[A] = tree match {
      case Nil => List[A]()
      case _ =>
        val child = for { t <- tree; child <- next(t) } yield child
        val values: List[A] = acc ::: flatten(tree)
        if(child != Nil) loop(child, values) else values
    }

    loop(List(this), List[A]())
  }
}

case class Leaf[A](value: A) extends Tree[A]
case class Node[A](value: A, left: Tree[A], right: Tree[A]) extends Tree[A]
case object Empty extends Tree[Nothing]

object Demo extends App {
  val t1 = Node(1, Node(2, Leaf(4), Leaf(5)), Node(3, Leaf(6), Leaf(7)))                                          // Int Tree
  val t2 = Node('F, Node('B, Leaf('A), Node('D, Leaf('C), Leaf('E))), Node('G, Empty, Node('I, Leaf('H), Empty))) //'Symbol Tree

  def print[T](t: Tree[T]) = {
    println("--------------------------------------------------------")
    println(t)
    println("========================================================")
    println(s"size\t\t${t.size}")
    println(s"height\t\t${t.height}")
    println(s"leafCount\t${t.leafCount}")
    println(s"leaves\t\t${t.leaves}")
    println(s"inOrder \t${t.inOrder}")
    println(s"preOrder\t${t.preOrder}")
    println(s"postOrder\t${t.postOrder}")
    println(s"levelOrder\t${t.levelOrder}")
  }

  print(t1)
  print(t2)
}

Demo.main(args)