Skip to content

Instantly share code, notes, and snippets.

@stew
Created February 4, 2016 06:59
Show Gist options
  • Save stew/905a6429462d73af58cf to your computer and use it in GitHub Desktop.
Save stew/905a6429462d73af58cf to your computer and use it in GitHub Desktop.
package dogs
import Predef._
import dogs.Order.{GT, EQ, LT, Ordering}
import dogs.std.intOrder
import scala.annotation.tailrec
import scala.math
sealed abstract class BinaryTree[A] {
import BinaryTree._
val size: Int
val height: Int
final def balanceFactor: Ordering = this match {
case BTNil() => EQ
case Branch(_,l,r) => Order[Int].apply(l.height,r.height)
}
def isEmpty: Boolean
def toLst(): List[A] = this match {
case BTNil() => El[A]
case Branch(a, l, r) => l.toLst ::: (a :: r.toLst)
}
def inorder() = toLst()
def min: Option[A] = {
@tailrec def loop(sub: BinaryTree[A], x: A): A = sub match {
case BTNil() => x
case Branch(a, l, _) => loop(l, a)
}
this match {
case BTNil() => None()
case Branch(a, l, _) => Some(loop(l, a))
}
}
def max: Option[A] = {
@tailrec def loop(sub: BinaryTree[A], x: A): A = sub match {
case BTNil() => x
case Branch(a, _, r) => loop(r, a)
}
this match {
case BTNil() => None()
case Branch(a, _, r) => Some(loop(r, a))
}
}
def contains(x: A)(implicit order: Order[A]): Boolean = this match {
case BTNil() => false
case Branch(a, l, r) => order.compare(x, a) match {
case LT => l.contains(x)
case EQ => true
case GT => r.contains(x)
}
}
def add(x: A)(implicit order: Order[A]): Branch[A] = {
def insert(x: A, order: Order[A]): Branch[A] = this match {
case BTNil() => Branch(x, BinaryTree.empty, BinaryTree.empty)
case branch @ Branch(a, l, r) => order.compare(x, a) match {
case LT => Branch(a, l.add(x)(order), r)
case EQ => branch
case GT => Branch(a, l, r.add(x)(order))
}
}
insert(x, order).balance
}
def +(x: A)(implicit order: Order[A]): BinaryTree[A] = add(x)
def remove(x: A)(implicit order: Order[A]): BinaryTree[A] = {
def del(x:A)(implicit order: Order[A]): BinaryTree[A] = this match {
case BTNil() => BinaryTree.empty
case Branch(a, l, r) => order.compare(x, a) match {
case LT => Branch(a, l.remove(x), r)
case GT => Branch(a, l, r.remove(x))
case EQ => r.min match {
case None() => l
case Some(v) => Branch(v,l,r.remove(v))
}
}
}
del(x) match {
case BTNil() => BTNil()
case b @ Branch(_,_,_) => b.balance
}
}
def join(another: BinaryTree[A])(implicit order: Order[A]) = {
@tailrec def build(sub: BinaryTree[A], xs: List[A]): BinaryTree[A] = xs match {
case El() => sub
case Nel(h, t) => build(sub + h, t)
}
another match {
case BTNil() => this
case _ => build(this, another.toLst())
}
}
def ++(another: BinaryTree[A])(implicit order: Order[A]) = join(another)
}
object BinaryTree {
def apply[A](a: A): BinaryTree[A] = Branch(a,empty,empty)
def empty[A]: BinaryTree[A] = BTNil()
trait Balance
case object L extends Balance
case object M extends Balance
case object R extends Balance
private[dogs] case class Branch[A] private[dogs](value: A,
left: BinaryTree[A],
right: BinaryTree[A]) extends BinaryTree[A] {
val size = left.size + right.size + 1
val height = java.lang.Math.max(left.height, right.height) + 1
override def isEmpty: Boolean = false
private[dogs] def balance: Branch[A] =
this.balanceFactor match {
case EQ => this
case GT => rotate(false)
case LT => rotate(true)
}
private[dogs] def rotate(rotateLeft: Boolean): Branch[A] =
if(rotateLeft)
right match {
case BTNil() => this
case Branch(rv,rl,rr) =>
Branch(rv,Branch(value,left,rl),rr)
}
else
left match {
case BTNil() => this
case Branch(lv,ll,lr) =>
Branch(lv,ll, Branch(value,lr,right))
}
}
private[dogs] case object BTNil extends BinaryTree[Nothing] {
override def isEmpty: Boolean = true
def apply[A](): BinaryTree[A] = this.asInstanceOf[BinaryTree[A]]
def unapply[A](a: BinaryTree[A]): Boolean = a.isEmpty
override val size: Int = 0
override val height: Int = 0
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment