Skip to content

Instantly share code, notes, and snippets.

@vkostyukov
Created September 2, 2013 01:38
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save vkostyukov/6408540 to your computer and use it in GitHub Desktop.
Save vkostyukov/6408540 to your computer and use it in GitHub Desktop.
A BST for PFDS in Scala (NSK usergroup)
import scala.annotation.tailrec
abstract sealed class Tree[+A] {
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)
*/
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](x: B)(implicit ord: Ordering[B]): Tree[B] = {
import ord._
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](x: B)(implicit ord: Ordering[B]): Tree[B] = {
import ord._
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