Data Structures in ReasonML: #6 AVL Tree

AvlTree.re

/* AVL Tree */
/* A balanced Binary Search Tree */
exception Undefined;

module type AvlTree = {
  type t = int;
  type height = int;
  type tree('t) =
    | Leaf
    | Node(height, tree('t), 't, tree('t));
  let empty: tree('t);
  let getHeight: (tree('t), tree('t)) => int;
  let getBalance: (tree('t), tree('t)) => int;
  let getTreeHeight: tree('t) => int;
  let rotateRight: tree('t) => tree('t);
  let rotateLeft: tree('t) => tree('t);
  let insert: (t, tree(t)) => tree(t);
  let isEmpty: tree(t) => bool;
  let contains: (t, tree(t)) => bool;
  let delete: (t, tree(t)) => tree(t);
  let isBalanced: tree('t) => bool;
};

module AvlTree: AvlTree = {
  type t = int; /* For now implement as int */
  type height = int;
  type tree('t) =
    | Leaf
    | Node(height, tree('t), 't, tree('t));
  let empty = Leaf;
  /* Return height for given tree */
  let getTreeHeight = tree =>
    switch (tree) {
    | Leaf => 0
    | Node(height, _, _, _) => height
    };
  /* Return the next height based on left and right tree height */
  let getHeight = (leftTree, rightTree) =>
    max(getTreeHeight(leftTree), getTreeHeight(rightTree)) + 1;
  /* Return balance for left and right node */
  let getBalance = (leftTree, rightTree) =>
    getTreeHeight(leftTree) - getTreeHeight(rightTree);
  /* Helper for checking if the tree is actually balanced */
  let rec isBalanced = tree =>
    switch (tree) {
    | Leaf => true
    | Node(height, leftTree, _, rightTree) =>
      height == getHeight(leftTree, rightTree)
      && abs(getBalance(leftTree, rightTree)) < 2
      && isBalanced(leftTree)
      && isBalanced(rightTree)
    };
  /* Check if value exists. We could rename to ~mem */
  let rec contains = (value, tree) =>
    switch (tree) {
    | Leaf => false
    | Node(_, leftTree, currentValue, rightTree) =>
      switch (value - currentValue) {
      | diff when diff < 0 => contains(value, leftTree)
      | diff when diff > 0 => contains(value, rightTree)
      | _ => true
      }
    };
  /* Rotate tree to left */
  let rotateLeft = tree =>
    switch (tree) {
    | Node(_, leftTree1, value1, Node(_, leftTree2, value2, rightTree2)) =>
      let leftTree =
        Node(getHeight(leftTree1, leftTree2), leftTree1, value1, leftTree2);
      Node(getHeight(leftTree, rightTree2), leftTree, value2, rightTree2);
    | _ => raise(Undefined)
    };
  /* Rotate tree to right */
  let rotateRight = tree =>
    switch (tree) {
    | Node(_, Node(_, leftTree2, value2, rightTree2), value1, rightTree1) =>
      let rightTree =
        Node(
          getHeight(rightTree2, rightTree1),
          rightTree2,
          value1,
          rightTree1,
        );
      Node(getHeight(leftTree2, rightTree), leftTree2, value2, rightTree);
    | _ => raise(Undefined)
    };
  let rec insert = (value, tree) =>
    switch (tree) {
    | Leaf => Node(1, Leaf, value, Leaf)
    | Node(_, leftTree, currentValue, rightTree) as tree =>
      switch (value - currentValue) {
      | diff when diff < 0 =>
        switch (insert(value, leftTree)) {
        | Node(_, lt, _, rt) as leftTree =>
          if (abs(getBalance(leftTree, rightTree)) < 2) {
            Node(
              getHeight(leftTree, rightTree),
              leftTree,
              currentValue,
              rightTree,
            );
          } else {
            let leftTree =
              if (getBalance(lt, rt) < 0) {
                rotateLeft(leftTree);
              } else {
                leftTree;
              };
            rotateRight(
              Node(
                getHeight(leftTree, rightTree),
                leftTree,
                currentValue,
                rightTree,
              ),
            );
          }
        | _ => raise(Undefined)
        }
      | diff when diff > 0 =>
        switch (insert(value, rightTree)) {
        | Node(_, lt, _, rt) as rightTree =>
          if (abs(getBalance(leftTree, rightTree)) < 2) {
            Node(
              getHeight(leftTree, rightTree),
              leftTree,
              currentValue,
              rightTree,
            );
          } else {
            let rightTree =
              if (getBalance(lt, rt) > 0) {
                rotateRight(rightTree);
              } else {
                rightTree;
              };
            rotateLeft(
              Node(
                getHeight(leftTree, rightTree),
                leftTree,
                currentValue,
                rightTree,
              ),
            );
          }
        | _ => raise(Undefined)
        }
      | _ => tree
      }
    };
  let isEmpty = tree => tree === Leaf;
  /* Return minimum value for given tree by going through the left side of the tree */
  let rec minimumValue = tree =>
    switch (tree) {
    | Node(_, Leaf, v, _) => v
    | Node(_, left, _, _) => minimumValue(left)
    | Leaf => raise(Undefined)
    };
  let rec delete = (value, tree) =>
    contains(value, tree) ?
      switch (tree) {
      | Leaf => Leaf
      | Node(_, leftTree, currentValue, rightTree) =>
        switch (value - currentValue) {
        | diff when diff < 0 =>
          let deletedLeftNode = delete(value, leftTree);
          let height = getHeight(deletedLeftNode, rightTree);
          let node = Node(height, deletedLeftNode, currentValue, rightTree);
          if (abs(getBalance(deletedLeftNode, rightTree)) > 1) {
            rotateLeft(node);
          } else {
            node;
          };
        | diff when diff > 0 =>
          let deletedRightNode = delete(value, rightTree);
          let height = getHeight(leftTree, deletedRightNode);
          let node = Node(height, leftTree, currentValue, deletedRightNode);
          if (abs(getBalance(leftTree, deletedRightNode)) > 1) {
            rotateLeft(node);
          } else {
            node;
          };
        | _ =>
          switch (leftTree, rightTree) {
          | (Leaf, Leaf) => Leaf
          | (Leaf, _) => rightTree
          | (_, Leaf) => leftTree
          | (_, _) =>
            let minimumValue = minimumValue(rightTree);
            let deletedTree = delete(minimumValue, rightTree);
            let node =
              Node(
                getHeight(leftTree, deletedTree),
                leftTree,
                minimumValue,
                deletedTree,
              );
            if (getBalance(leftTree, deletedTree) > 1) {
              rotateRight(node);
            } else {
              node;
            };
          }
        }
      } :
      tree;
};

let run = () => {
  let t = AvlTree.empty;
  Js.log("Insert values");
  let t = AvlTree.insert(1, t);
  let t = AvlTree.insert(12, t);
  let t = AvlTree.insert(3, t);
  let t = AvlTree.insert(40, t);
  let t = AvlTree.insert(20, t);
  let t = AvlTree.insert(7, t);
  let t = AvlTree.insert(27, t);
  let t = AvlTree.insert(37, t);
  let t = AvlTree.insert(97, t);
  let t = AvlTree.insert(57, t);
  let t = AvlTree.insert(77, t);
  let t = AvlTree.insert(17, t);
  let t = AvlTree.insert(27, t);
  Js.log2("Is balanced? ", AvlTree.isBalanced(t));
  Js.log("test");
  Js.log2("contains 27?", AvlTree.contains(27, t));
  Js.log("remove 27");
  let t = AvlTree.delete(27, t);
  Js.log2("contains 27?", AvlTree.contains(27, t));
  Js.log2("Is balanced? ", AvlTree.isBalanced(t));
  Js.log2("contains 3?", AvlTree.contains(3, t));
  Js.log("remove 3");
  let t = AvlTree.delete(3, t);
  Js.log2("contains 3?", AvlTree.contains(3, t));
  Js.log2("contains 77?", AvlTree.contains(77, t));
  Js.log("remove 77");
  let t = AvlTree.delete(77, t);
  Js.log2("contains 77?", AvlTree.contains(77, t));
  Js.log2("contains 97?", AvlTree.contains(97, t));
  Js.log("remove 97");
  let t = AvlTree.delete(97, t);
  Js.log2("contains 97?", AvlTree.contains(97, t));
  Js.log2("contains 1?", AvlTree.contains(1, t));
  Js.log("remove 1");
  let t = AvlTree.delete(1, t);
  Js.log2("contains 1?", AvlTree.contains(1, t));
  Js.log2("Is balanced? ", AvlTree.isBalanced(t));
};

run();