AVLTrees (implemented in TypeScript, D)

avl_tree.d

module avl;

import std.stdio;
import std.typecons;

string string_rep(string s, size_t n) {
  string ret;
  foreach (_; 0 .. n) {
    ret ~= s;
  }
  return ret;
}

class AVLTree(K, V) {
  private AVLNode!(K, V) root;

  private class AVLNode(K, V) {
    public K key;
    public V value;
    private int height, size;
    private AVLNode!(K, V) left;
    private AVLNode!(K, V) right;

    this(K key, V value) {
      this.key = key;
      this.value = value;
      this.height = 1;
      this.size = 1;
    }
  }

  public Nullable!V find(K key) {
    AVLNode!(K, V) ret = AVLTree!(K, V).find(this.root, key);

    if (ret is null) {
      return typeof(return).init;
    } else {
      return nullable(ret.value);
    }
  }

  public bool exists(K key) {
    return AVLTree!(K, V).find(this.root, key) !is null;
  }

  private static AVLNode!(K, V) find(AVLNode!(K, V) t, K key) {
    if (t is null) {
      return null;
    }
    if (key == t.key) {
      return t;
    } else if (key < t.key) {
      return find(t.left, key);
    } else {
      return find(t.right, key);
    }
  }

  public void insert(K key, V value) {
    this.root = AVLTree!(K, V).insert(root, new AVLNode!(K, V)(key, value));
  }

  private static AVLNode!(K, V) insert(AVLNode!(K, V) t, AVLNode!(K, V) x) {
    if (t is null) {
      return x;
    }

    if (x.key == t.key) {
      t.value = x.value;
    } else if (x.key < t.key) {
      t.left = AVLTree!(K, V).insert(t.left, x);
    } else {
      t.right = AVLTree!(K, V).insert(t.right, x);
    }
    t.size += 1;

    return AVLTree!(K, V).balance(t);
  }

  private static int sz(AVLNode!(K, V) t) {
    return t !is null ? t.size : 0;
  }

  private static int ht(AVLNode!(K, V) t) {
    return t !is null ? t.height : 0;
  }

  enum LR {
    L,
    R
  }

  private static AVLNode!(K, V) get_child_by_LR(AVLNode!(K, V) t, LR lr) {
    final switch (lr) with (LR) {
    case L:
      return t.left;
    case R:
      return t.right;
    }
  }

  private static void set_child_by_LR(AVLNode!(K, V) dst, LR lr, AVLNode!(K, V) src) {
    final switch (lr) with (LR) {
    case L:
      dst.left = src;
      break;
    case R:
      dst.right = src;
      break;
    }
  }

  private static AVLNode!(K, V) rotate(AVLNode!(K, V) t, LR l, LR r) {
    AVLNode!(K, V) s = get_child_by_LR(t, r);
    AVLTree!(K, V).set_child_by_LR(t, r, get_child_by_LR(s, l));
    AVLTree!(K, V).set_child_by_LR(s, l, AVLTree!(K, V).balance(t));

    if (t !is null) {
      t.size = sz(t.left) + sz(t.right) + 1;
    }
    if (s !is null) {
      s.size = sz(s.left) + sz(s.right) + 1;
    }

    return AVLTree!(K, V).balance(s);
  }

  private static T max(T)(T a, T b) {
    return a > b ? a : b;
  }

  private static AVLNode!(K, V) balance(AVLNode!(K, V) t) {
    if (ht(t.right) - ht(t.left) < -1) {
      if (ht(t.left.right) - ht(t.left.left) > 0) {
        t.left = AVLTree!(K, V).rotate(t.left, LR.L, LR.R);
      }
      return AVLTree!(K, V).rotate(t, LR.R, LR.L);
    }

    if (ht(t.left) - ht(t.right) < -1) {
      if (ht(t.right.left) - ht(t.right.right) > 0) {
        t.right = AVLTree!(K, V).rotate(t.right, LR.R, LR.L);
      }
      return AVLTree!(K, V).rotate(t, LR.L, LR.R);
    }

    if (t !is null) {
      t.height = AVLTree!(K, V).max(ht(t.left), ht(t.right)) + 1;
      t.size = sz(t.left) + sz(t.right) + 1;
    }
    return t;
  }

  public static void print_node(AVLNode!(K, V) node, size_t depth = 0) {
    if (node !is null) {
      print_node(node.left, depth + 1);
      writefln("%s <%s:%s>", string_rep("    ", depth), node.key, node.value);
      print_node(node.right, depth + 1);
    }
  }

  public void print_tree() {
    print_node(this.root, 0);
  }
}

void main() {
  import std.random;

  AVLTree!(int, int) tree = new AVLTree!(int, int);

  Mt19937 mt;
  mt.seed(139);

  enum SIZE = 20000;
  enum MIN_NUM = 0;
  enum MAX_NUM = 10000000;

  writeln("Source vector size : ", SIZE);
  writeln("Minimum element of source element vector : ", MIN_NUM);
  writeln("Maximum element of source element vector : ", MAX_NUM);

  int[] vs = new int[SIZE];
  
  enum RANDOM = true;

  foreach (i; 0 .. SIZE) {
    vs[i] = RANDOM ? uniform(MIN_NUM, MAX_NUM, mt) : i;
  }

  import std.datetime.stopwatch;

  StopWatch sw_all, sw_insert, sw_search;
  writeln("insert start...");

  sw_all.start;
  sw_insert.start;
  foreach (v; vs) {
    tree.insert(v, v);
  }
  sw_insert.stop;

  writeln("insert completed... ", sw_insert.peek);

  writeln("search start...");
  sw_search.start;
  foreach (v; vs) {
    tree.find(v);
  }
  sw_search.stop;
  sw_all.stop;
  writeln("search completed... ", sw_search.peek);

  writeln("[summary] ------------------------");
  writeln("total  : ", sw_all.peek);
  writeln("insert : ", sw_insert.peek);
  writeln("search : ", sw_search.peek);

  //tree.print_tree();
}

avl_tree.ts

import { Option, none, some } from './option';

function string_rep(s: string, n: number) {
  let ret = "";
  for (let i = 0; i < n; i++) {
    ret += s;
  }
  return ret;
}

class AVLNode<K, V> {
  key: K;
  value: V;
  height: number;
  size: number;
  left: AVLNode<K, V>;
  right: AVLNode<K, V>;

  constructor(key: K, value: V) {
    this.key = key;
    this.value = value;
    this.height = 1;
    this.size = 1;
  }
}

class LR {
  static L = 0;
  static R = 1;
}

function max<T>(a: T, b: T): T {
  return a > b ? a : b;
}

interface Less {
  type: 'Less';
}
interface Equal {
  type: 'Equal';
}
interface Greater {
  type: 'Greater';
}
export type CompareResult =
  | Less
  | Equal
  | Greater;

export function less(): CompareResult {
  return { type: 'Less' };
}
export function equal(): CompareResult {
  return { type: 'Equal' };
}
export function greater(): CompareResult {
  return { type: 'Greater' };
}

export type CompareFunction<T> = (a: T, b: T) => CompareResult;
export type PrinterFunction<T> = (e: T) => string;

export class AVLTree<K, V> {
  protected root: AVLNode<K, V>;

  public size(): number {
    if (this.root === undefined) {
      return 0;
    } else {
      return AVLTree.sz(this.root);
    }
  }

  public find(key: K): Option<V> {
    let ret: Option<AVLNode<K, V>> = AVLTree.find<K, V>(this.root, key);

    if (ret.type === 'None') {
      return none<V>();
    } else {
      return some(ret.value.value);
    }
  }

  public exists(key: K): boolean {
    return AVLTree.find<K, V>(this.root, key).type !== 'None';
  }

  protected static find<K, V>(t: AVLNode<K, V>, key: K): Option<AVLNode<K, V>> {
    if (t === undefined) {
      return none();
    }
    if (key === t.key) {
      return some(t);
    } else if (key < t.key) {
      return AVLTree.find(t.left, key);
    } else {
      return AVLTree.find(t.right, key);
    }
  }

  public findWithCompareFunction(key: K, compare: CompareFunction<K>): Option<V> {
    let ret: Option<AVLNode<K, V>> = AVLTree.findWithCompareFunction<K, V>(this.root, key, compare);

    if (ret.type === 'None') {
      return none<V>();
    } else {
      return some(ret.value.value);
    }
  }

  protected static findWithCompareFunction<K, V>(t: AVLNode<K, V>, key: K, compare: CompareFunction<K>): Option<AVLNode<K, V>> {
    if (t === undefined) {
      return none();
    }

    let cmp = compare(key, t.key);

    switch (cmp.type) {
      case 'Equal': return some(t);
      case 'Less': return AVLTree.findWithCompareFunction(t.left, key, compare);
      case 'Greater': return AVLTree.findWithCompareFunction(t.right, key, compare);
    }
  }

  public insert(key: K, value: V): void {
    this.root = AVLTree.insert(this.root, new AVLNode<K, V>(key, value));
  }

  protected static insert<K, V>(t: AVLNode<K, V>, x: AVLNode<K, V>): AVLNode<K, V> {
    if (t === undefined) {
      return x;
    }

    if (x.key === t.key) {
      t.value = x.value;
    } else if (x.key < t.key) {
      t.left = AVLTree.insert(t.left, x);
    } else {
      t.right = AVLTree.insert(t.right, x);
    }
    t.size += 1;

    return AVLTree.balance(t);
  }

  public insertWithCompareFunction(key: K, value: V, compare: CompareFunction<K>): void {
    this.root = AVLTree.insertWithCompareFuncion(this.root, new AVLNode<K, V>(key, value), compare);
  }

  protected static insertWithCompareFuncion<K, V>(t: AVLNode<K, V>, x: AVLNode<K, V>, compare: CompareFunction<K>): AVLNode<K, V> {
    if (t === undefined) {
      return x;
    }

    let cmp = compare(x.key, t.key);

    switch (cmp.type) {
      case 'Equal':
        t.value = x.value; break;
      case 'Less':
        t.left = AVLTree.insertWithCompareFuncion(t.left, x, compare); break;
      case 'Greater':
        t.right = AVLTree.insertWithCompareFuncion(t.right, x, compare); break;
    }

    t.size += 1;

    return AVLTree.balance(t);
  }

  protected static sz<K, V>(t: AVLNode<K, V>): number {
    return t !== undefined ? t.size : 0;
  }

  protected static ht<K, V>(t: AVLNode<K, V>): number {
    return t !== undefined ? t.height : 0;
  }

  protected static get_child_by_LR<K, V>(t: AVLNode<K, V>, lr: LR): AVLNode<K, V> {
    switch (lr) {
      case LR.L:
        return t.left;
      case LR.R:
        return t.right;
    }
  }

  protected static set_child_by_LR<K, V>(dst: AVLNode<K, V>, lr: LR, src: AVLNode<K, V>): void {
    switch (lr) {
      case LR.L:
        dst.left = src; break;
      case LR.R:
        dst.right = src; break;
    }
  }

  protected static rotate<K, V>(t: AVLNode<K, V>, l: LR, r: LR): AVLNode<K, V> {
    let s: AVLNode<K, V> = AVLTree.get_child_by_LR(t, r);
    AVLTree.set_child_by_LR(t, r, AVLTree.get_child_by_LR(s, l));
    AVLTree.set_child_by_LR(s, l, AVLTree.balance(t));

    if (t !== undefined) {
      t.size = AVLTree.sz(t.left) + AVLTree.sz(t.right) + 1;
    }
    if (s !== undefined) {
      s.size = AVLTree.sz(s.left) + AVLTree.sz(s.right) + 1;
    }

    return AVLTree.balance(s);
  }

  protected static balance<K, V>(t: AVLNode<K, V>): AVLNode<K, V> {
    if (AVLTree.ht(t.right) - AVLTree.ht(t.left) < -1) {
      if (AVLTree.ht(t.left.right) - AVLTree.ht(t.left.left) > 0) {
        t.left = AVLTree.rotate(t.left, LR.L, LR.R);
      }
      return AVLTree.rotate(t, LR.R, LR.L);
    }

    if (AVLTree.ht(t.left) - AVLTree.ht(t.right) < -1) {
      if (AVLTree.ht(t.right.left) - AVLTree.ht(t.right.right) > 0) {
        t.right = AVLTree.rotate(t.right, LR.R, LR.L);
      }
      return AVLTree.rotate(t, LR.L, LR.R);
    }

    if (t !== undefined) {
      t.height = max(AVLTree.ht(t.left), AVLTree.ht(t.right)) + 1;
      t.size = AVLTree.sz(t.left) + AVLTree.sz(t.right) + 1;
    }
    return t;
  }

  public delete(key: K) {
    this.root = AVLTree.delete(this.root, key);
  }

  public deleteWithCompareFunction(key: K, compare: CompareFunction<K>) {
    this.root = AVLTree.deleteWithCompareFunction(this.root, key, compare);
  }

  private static delete<K, V>(t: AVLNode<K, V>, x: K): AVLNode<K, V> {
    if (t === undefined) {
      return undefined;
    }
    if (x === t.key) {
      return AVLTree.move_down(t.left, t.right);
    }
    if (x < t.key) {
      t.left = AVLTree.delete(t.left, x);
    } else {
      t.right = AVLTree.delete(t.right, x);
    }
    t.size -= 1;
    return AVLTree.balance(t);
  }

  private static deleteWithCompareFunction<K, V>(t: AVLNode<K, V>, x: K, compare: CompareFunction<K>): AVLNode<K, V> {
    if (t === undefined) {
      return undefined;
    }

    let comp = compare(x, t.key);

    if (comp.type === 'Equal') {
      return AVLTree.move_down(t.left, t.right);
    }
    if (comp.type === 'Less') {
      t.left = AVLTree.deleteWithCompareFunction(t.left, x, compare);
    } else {
      t.right = AVLTree.deleteWithCompareFunction(t.right, x, compare);
    }
    t.size -= 1;
    return AVLTree.balance(t);
  }

  protected static move_down<K, V>(t: AVLNode<K, V>, rhs: AVLNode<K, V>): AVLNode<K, V> {
    if (t === undefined) {
      return rhs;
    }

    t.right = AVLTree.move_down(t.right, rhs);
    return AVLTree.balance(t);
  }

  protected static printNode<K, V>(node: AVLNode<K, V>, depth: number = 0): void {
    if (node !== undefined) {
      AVLTree.printNode(node.left, depth + 1);
      console.log(`${string_rep("    ", depth)} <${node.key}, ${node.value}>`);
      AVLTree.printNode(node.right, depth + 1);
    }
  }

  protected printTree(): void {
    AVLTree.printNode(this.root, 0);
  }

  public static printNodeWithPrinterFunction<K, V>(
    node: AVLNode<K, V>, depth: number = 0,
    key_printer: PrinterFunction<K>, value_printer: PrinterFunction<V>): void {
    if (node !== undefined) {
      AVLTree.printNodeWithPrinterFunction(node.left, depth + 1, key_printer, value_printer);
      console.log(`${string_rep("    ", depth)} <${key_printer(node.key)}, ${value_printer(node.value)}>`);
      AVLTree.printNodeWithPrinterFunction(node.right, depth + 1, key_printer, value_printer);
    }
  }

  public printTreeWithPrinterFunction(key_printer: PrinterFunction<K>, value_printer: PrinterFunction<V>): void {
    AVLTree.printNodeWithPrinterFunction(this.root, 0, key_printer, value_printer);
  }

  protected static recursiveCollector<K, V, E>(node: AVLNode<K, V>, selector: (node: AVLNode<K, V>) => E): E[] {
    if (node === undefined) {
      return [];
    }

    let ret: E[] = [];

    ret.push(selector(node));

    if (node.left !== undefined) {
      ret = ret.concat(AVLTree.recursiveCollector(node.left, selector));
    }

    if (node.right !== undefined) {
      ret = ret.concat(AVLTree.recursiveCollector(node.right, selector));
    }

    return ret
  }

  protected static keys<K, V>(node: AVLNode<K, V>): K[] {
    return AVLTree.recursiveCollector(node, (node: AVLNode<K, V>) => node.key);
  }

  public keys(): K[] {
    return AVLTree.keys(this.root);
  }

  protected static values<K, V>(node: AVLNode<K, V>): V[] {
    return AVLTree.recursiveCollector(node, (node: AVLNode<K, V>) => node.value);
  }

  public values(): V[] {
    return AVLTree.values(this.root);
  }
}

function getRandomInt(max: number) {
  return Math.floor(Math.random() * Math.floor(max));
}

let now = require('performance-now');

export class StopWatch {
  private start_tm: number = 0;
  private end_tm: number = 0;

  start(): void {
    this.start_tm = now();
  }

  stop(): void {
    this.end_tm = now();
  }

  peek(): number {
    let t = now();
    return t - this.start_tm;
  }

  diff(): number {
    return this.end_tm - this.start_tm;
  }
}

class Container {
  key: number;
  value: number;

  constructor(v: number) {
    this.key = v;
    this.value = v;
  }
}

export function avl_test_main(random = true) {
  let tree = new AVLTree<number, number>();

  const SIZE = 20000;
  const MIN_NUM = 0;
  const MAX_NUM = 10000000;

  console.log("Source vector size : ", SIZE);
  console.log("Minimum element of source element vector : ", MIN_NUM);
  console.log("Maximum element of source element vector : ", MAX_NUM);

  let vs = new Array<number>(SIZE);

  for (let i = 0; i < SIZE; i++) {
    vs[i] = random ? getRandomInt(MAX_NUM) : i;
  }

  let sw_all = new StopWatch();
  let sw_insert = new StopWatch();
  let sw_search = new StopWatch();
  console.log("insert start...");

  sw_all.start();
  sw_insert.start();

  vs.forEach(v => tree.insert(v, v));
  sw_insert.stop();

  console.log("insert completed... ", sw_insert.diff());

  console.log("search start...");
  sw_search.start();
  vs.forEach(v => tree.find(v));

  sw_search.stop();
  sw_all.stop();
  console.log("search completed... ", sw_search.diff());

  console.log("[summary] ------------------------");
  console.log("total  : ", sw_all.diff());
  console.log("insert : ", sw_insert.diff());
  console.log("search : ", sw_search.diff());
}

option.ts

export interface Some<T> {
  type: 'Some';
  value: T;
}
export interface None {
  type: 'None';
}
export type Option<T> = Some<T> | None;

export function some<T>(v: T): Option<T> {
  return {
    type: 'Some',
    value: v
  };
}

export function none<T>(): Option<T> {
  return {
    type: 'None'
  };
}

export function isNone<T>(obj: Option<T>): boolean {
  return obj.type === 'None';
}

export function isSome<T>(obj: Option<T>): boolean {
  return obj.type === 'Some';
}

export function map<T, U>(obj: Option<T>, f: (obj: T) => U): Option<U> {
  if (obj.type === 'Some') {
    return some(f(obj.value));
  } else {
    return none<U>();
  }
}

export function bind<T, U>(obj: Option<T>, f: (obj: T) => Option<U>): Option<U> {
  if (obj.type === 'Some') {
    return f(obj.value);
  } else {
    return none<U>();
  }
}

export function _default<T>(obj: Option<T>, _default: T): T {
  if (obj.type === 'Some') {
    return obj.value;
  } else {
    return _default;
  }
}

export function may<T>(obj: Option<T>, f: (obj: T) => void): void {
  if (obj.type === 'Some') {
    f(obj.value);
  }
}

export function map_default<T, U>(obj: Option<T>, f: (obj: T) => U, _default: U): U {
  if (obj.type === 'Some') {
    return f(obj.value);
  } else {
    return _default;
  }
}