Skip to content

Instantly share code, notes, and snippets.

@cypok
Forked from vkostyukov/Tree.scala
Last active December 22, 2015 03:18
Show Gist options
  • Save cypok/6408880 to your computer and use it in GitHub Desktop.
Save cypok/6408880 to your computer and use it in GitHub Desktop.
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment