SplayTree

SplayTree.ts

/**
 * SplayTree is weighted binary search tree which is based on SplayTree.
 *  - https://en.wikipedia.org/wiki/Splay_tree
 *
 * Original paper on SplayTree:
 *  - https://www.cs.cmu.edu/~sleator/papers/self-adjusting.pdf
 */
export abstract class SplayNode<V> {
  protected value: V;

  private left: SplayNode<V> = null;
  private right: SplayNode<V> = null;
  private parent: SplayNode<V> = null;
  private weight: number;

  constructor(value: V) {
    this.value = value;
    this.initWeight();
  }

  public abstract getLength();

  public getValue(): V {
    return this.value;
  }

  public getLeftWeight(): number {
    return !this.hasLeft() ? 0 : this.left.getWeight();
  }

  public getRightWeight(): number {
    return !this.hasRight() ? 0 : this.right.getWeight();
  }

  public getWeight(): number {
    return this.weight;
  }

  public getLeft(): SplayNode<V> {
    return this.left;
  }

  public getRight(): SplayNode<V> {
    return this.right;
  }

  public setRight(right: SplayNode<V>): void {
    this.right = right;
  }

  public hasLeft(): boolean {
    return this.left != null;
  }

  public hasRight(): boolean {
    return this.right != null;
  }

  public hasParent(): boolean {
    return this.parent != null;
  }

  public setParent(parent: SplayNode<V>): void {
    this.parent = parent;
  }

  public setLeft(left: SplayNode<V>): void {
    this.left = left;
  }

  public getParent(): SplayNode<V> {
    return this.parent;
  }

  public increaseWeight(weight: number): void {
    this.weight += weight;
  }

  public initWeight(): void {
    this.weight = this.getLength();
  }
}

export class SplayTree<V> {
  private root: SplayNode<V>;

  constructor(root: SplayNode<V>) {
    this.root = root;
  }

  public find(pos: number): [SplayNode<V>, number] {
    let node = this.root;
    while (true) {
      if (node.hasLeft() && pos <= node.getLeftWeight()) {
        node = node.getLeft();
      } else if (node.hasRight() && node.getLeftWeight() + node.getLength() < pos) {
        pos -= node.getLeftWeight() + node.getLength();
        node = node.getRight();
      } else {
        pos -= node.getLeftWeight();
        break;
      }
    }

    return [node, pos];
  }

  /**
   * Find the index of the given node in BST.
   * @param node the given node
   * @return the index of given node
   */
  public indexOf(node: SplayNode<V>): number {
    if (node === null) {
      return -1;
    }

    let index = 0;
    let current = node;
    let prev = null;
    while (current !== null) {
      if (prev === null || prev == current.getRight()) {
        index += current.getLength() + (current.hasLeft() ? current.getLeftWeight() : 0);
      }
      prev = current;
      current = current.getParent();
    }
    return index - node.getLength();
  }

  public getRoot(): SplayNode<V> {
    return this.root;
  }

  public insertAfter(targetOrNewNode: SplayNode<V>, newNode?: SplayNode<V>): SplayNode<V> {
    const target = typeof newNode !== 'undefined' ? targetOrNewNode : this.root;
    const newOne = typeof newNode !== 'undefined' ? newNode : targetOrNewNode;

    this.splayNode(target);
    this.root = newOne;
    newOne.setRight(target.getRight());
    if (target.hasRight()) {
      target.getRight().setParent(newOne);
    }
    newOne.setLeft(target);
    target.setParent(newOne);
    target.setRight(null);
    this.updateSubtree(target);
    this.updateSubtree(newOne);

    return newOne;
  }

  public updateSubtree(node: SplayNode<V>): void {
    node.initWeight();

    if (node.hasLeft()) {
      node.increaseWeight(node.getLeftWeight());
    }
    if (node.hasRight()) {
      node.increaseWeight(node.getRightWeight());
    }
  }

  public splayNode(node: SplayNode<V>): void {
    if (node === null) {
      return;
    }

    while (true) {
      if (this.isLeftChild(node.getParent()) && this.isRightChild(node)) { // zig-zag
        this.rotateLeft(node);
        this.rotateRight(node);
      } else if (this.isRightChild(node.getParent()) && this.isLeftChild(node)) { // zig-zag
        this.rotateRight(node);
        this.rotateLeft(node);
      } else if (this.isLeftChild(node.getParent()) && this.isLeftChild(node)) { // zig-zig
        this.rotateRight(node.getParent());
        this.rotateRight(node);
      } else if (this.isRightChild(node.getParent()) && this.isRightChild(node)) { // zig-zig
        this.rotateLeft(node.getParent());
        this.rotateLeft(node);
      } else { // zig
        if (this.isLeftChild(node)) {
          this.rotateRight(node);
        } else if (this.isRightChild(node)) {
          this.rotateLeft(node);
        }
        return;
      }
    }
  }

  public rotateLeft(pivot: SplayNode<V>): void {
    const root = pivot.getParent();
    if (root.hasParent()) {
      if (root === root.getParent().getLeft()) {
        root.getParent().setLeft(pivot);
      } else {
        root.getParent().setRight(pivot);
      }
    } else {
      this.root = pivot;
    }
    pivot.setParent(root.getParent());

    root.setRight(pivot.getLeft());
    if (root.hasRight()) {
      root.getRight().setParent(root);
    }

    pivot.setLeft(root);
    pivot.getLeft().setParent(pivot);

    this.updateSubtree(root);
    this.updateSubtree(pivot);
  }

  public getAnnotatedString(): string {
    const stack = [];
    this.traverseInorder(this.root, stack);
    return stack.map((node) =>
      `[${node.getWeight()},${node.getLength()}:${node.getValue().toString()}]`
    ).join('');
  }

  private traverseInorder(node: SplayNode<V>, stack: Array<SplayNode<V>>): void {
    if (node === null) {
      return;
    }

    this.traverseInorder(node.getLeft(), stack);
    stack.push(node);
    this.traverseInorder(node.getRight(), stack);
  }

  private rotateRight(pivot: SplayNode<V>): void {
    const root = pivot.getParent();
    if (root.hasParent()) {
      if (root === root.getParent().getLeft()) {
        root.getParent().setLeft(pivot);
      } else {
        root.getParent().setRight(pivot);
      }
    } else {
      this.root = pivot;
    }
    pivot.setParent(root.getParent());

    root.setLeft(pivot.getRight());
    if (root.hasLeft()) {
      root.getLeft().setParent(root);
    }

    pivot.setRight(root);
    pivot.getRight().setParent(pivot);

    this.updateSubtree(root);
    this.updateSubtree(pivot);
  }

  private isLeftChild(node: SplayNode<V>): boolean {
    return node !== null && node.hasParent() && node.getParent().getLeft() === node;
  }

  private isRightChild(node: SplayNode<V>): boolean {
    return node !== null && node.hasParent() && node.getParent().getRight() === node;
  }
}