Tree.scala

import scala.annotation.tailrec

object TestRun extends App {
  val tree = Leaf.add(10).add(5).add(7)
  println(tree.valuesByBreadth mkString ", ")
}

abstract sealed class Tree[+A] {
  import Ordered._

  def value: A
  def left: Tree[A]
  def right: Tree[A]
  def isEmpty: Boolean
  def size: Int

  /**
   * Time - O(log n)
   * Space - O(log n)
   */
  def min: A = {
    @tailrec def loop(t: Tree[A], m: A): A =
      if (t.isEmpty) m else loop(t.left, t.value)

    if (isEmpty) fail("Empty tree.")
    else loop(left, value)
  }

  /**
   * Time - O(log n)
   * Space - O(log n)
   */
  def max: A = {
    @tailrec def loop(t: Tree[A], m: A): A =
      if (t.isEmpty) m else loop(t.right, t.value)

    if (isEmpty) fail("Empty tree.")
    else loop(right, value)
  }

  /**
   * Time - O(1)
   * Space - O(1)
   * TODO: why factory method is here?
   */
  def mkTree[B : Ordering](v: B, l: Tree[B] = Leaf, r: Tree[B] = Leaf): Tree[B] =
    Branch(v, l, r, l.size + r.size + 1)

  /**
   * Time - O(log n)
   * Space - O(log n)
   */
  def add[B >: A : Ordering](x: B): Tree[B] = {
    if (isEmpty) mkTree(x)
    else if (x < value) mkTree(value, left.add(x), right)
    else if (x > value) mkTree(value, left, right.add(x))
    else this
  }

  /**
   * Time - O(log n)
   * Space - O(log n)
   */
  def remove[B >: A : Ordering](x: B): Tree[B] = {
    if (isEmpty) fail("Can't find " + x + " in this tree.")
    else if (x < value) mkTree(value, left.remove(x), right)
    else if (x > value) mkTree(value, left, right.remove(x))
    else {
      if (left.isEmpty && right.isEmpty) Leaf // case 1
      else if (left.isEmpty) right            // case 2
      else if (right.isEmpty) left            // case 2
      else {                                  // case 3
        val succ: B = right.min
        mkTree(succ, left, right.remove(succ))
      }
    }
  }

  /**
   * Time - O(log n)
   * Space - O(log n)
   */
  def apply(n: Int): A =
    if (isEmpty) fail("Tree doesn't contain a " + n + "th element.")
    else if (n < left.size) left(n)
    else if (n > left.size) right(n - size - 1)
    else value

  /**
   * Time - O(n)
   * Space - O(log n)
   */
  def valuesByDepth: List[A] = {
    def loop(s: List[Tree[A]]): List[A] =
      if (s.isEmpty) Nil
      else if (s.head.isEmpty) loop(s.tail)
      else s.head.value :: loop(s.head.right :: s.head.left :: s.tail)

    loop(List(this))
  }

  /**
   * Time - O(n)
   * Space - O(log n)
   */
  def valuesByBreadth: List[A] = {
    import scala.collection.immutable.Queue
    def loop(q: Queue[Tree[A]]): List[A] =
      if (q.isEmpty) Nil
      else if (q.head.isEmpty) loop(q.tail)
      else q.head.value :: loop(q.tail :+ q.head.left :+ q.head.right)

    loop(Queue(this))
  }

  /**
   * Time - O(n)
   * Space - O(log n)
   */
  def invert[B >: A](implicit num: Numeric[B]): Tree[B] =
    if (isEmpty) Leaf
    else mkTree(num.negate(value), right.invert(num), left.invert(num))

  /**
   * Fails with message.
   */
  def fail(s: String): Nothing =
    throw new NoSuchElementException(s)
}

case object Leaf extends Tree[Nothing] {
  def value: Nothing = fail("Empty tree.")
  def left: Tree[Nothing] = fail("Empty tree.")
  def right: Tree[Nothing] = fail("Empty tree.")
  def size: Int = 0
  def isEmpty: Boolean = true
}

case class Branch[A : Ordering](value: A,
                               left: Tree[A],
                               right: Tree[A],
                               size: Int) extends Tree[A] {
  def isEmpty: Boolean = false
}