splaytree.scala

splaytree.scala

package splaytree;

private final class BinaryNode[K, V](
  var key : K, 
  var value : V){

  var left : BinaryNode[K, V] = null
  var right : BinaryNode[K, V] = null

  final def minKey : K = if (left != null) left.minKey else key
}

import scala.collection._

class SplayTree[K, V](implicit ord : Ordering[K]) extends mutable.Map[K, V] {
  import ord.compare;
  private[this] type Node = BinaryNode[K, V];

  private var root : Node = null;
  private var _size = 0;

  def size = _size;  

  private class IteratorImpl(start : Option[K], cap : Option[K]) extends Iterator[(K, V)]{
    var nextKey = start.getOrElse(firstKey)

    var _hasNext = true;

    def hasNext = _hasNext && (cap match { case None => true; case Some(x) => compare(nextKey, x) < 0})

    def next = {
      splay(nextKey);
      if (root.right == null) _hasNext = false;
      else nextKey = root.right.minKey;
      (root.key, root.value);
    }
  }

  def elements : Iterator[(K, V)] = 
    if (root == null) Iterator.empty
    else new IteratorImpl(None, None);

  def update(key : K, value : V) {
    if (root == null) {
      root = new BinaryNode(key, value);
      _size += 1;
      return;
    }
    splay(key);
    val c = compare(key, root.key);
    if (c == 0) {
      root.value = value;
      return;
    }
    _size += 1;
    val n = new BinaryNode(key, value);
    if (c < 0) {
      n.left = root.left;
      n.right = root;
      root.left = null;
    } else {
      n.right = root.right;
      n.left = root;
      root.right = null;
    }
    root = n;
  }

  def -=(key : K) {
    splay(key);
    if (compare(key, root.key) != 0) {
      return;
    } 
    _size -= 1;

    // Now delete the root
    if (root.left == null) {
        root = root.right;
    } else {
      val x = root.right;
      root = root.left;
      splay(key);
      root.right = x;
    }
  }

  def firstKey = {
    var x = root;
    if(root == null) error("empty map") 
    while(x.left != null) x = x.left;

    splay(x.key);

    x.key;
  }

  def lastKey  = {
    var x = root;
    if(root == null) error("empty map");
    while(x.right != null) x = x.right;
    splay(x.key);
    x.key
  }

  def get(key : K) : Option[V] = {
    if (root == null) return None;
    splay(key);
    if(compare(root.key, key) != 0) return None;
    Some(root.value);
  }

  def empty = root == null;

  private def splay(key : K) {
    val header = new BinaryNode(null, null).asInstanceOf[Node]; 
    var l, r = header;
    var t = root; 
    var y : Node = null;
    def doSplay {
      while(true){
        val c = compare(key, t.key);
        if (c < 0) {
          if (t.left == null) return;
          if (compare(key, t.left.key) < 0) {
            y = t.left;                            /* rotate right */
            t.left = y.right;
            y.right = t;
            t = y;
            if (t.left == null) return;
          }
          r.left = t;                                 /* link right */
          r = t;
          t = t.left;
        } else if (c > 0) {
          if (t.right == null) return;
          if (compare(key, t.right.key) > 0) {
            y = t.right;                            /* rotate left */
            t.right = y.left;
            y.left = t;
            t = y;
            if (t.right == null) return;
          }
          l.right = t;                                /* link left */
          l = t;
          t = t.right;
        } else return;
      }
    }

    doSplay;
    l.right = t.left;                                   /* assemble */
    r.left = t.right;
    t.left = header.right;
    t.right = header.left;
    root = t;
  }
}