Implementation of AA Tree

aa-tree.scala

package ds.tree.aa

import scala.annotation.tailrec

// https://en.wikipedia.org/wiki/AA_tree

sealed trait Tree

final case class Node(var key: Int, var level: Int, var left: Tree, var right: Tree) extends Tree

case object Nil extends Tree

object Tree {
  def make(): Tree = Nil

  def leaf(key: Int): Tree = Node(key, 1, Nil, Nil)

  def leaf(key: Int, left: Tree, right: Tree): Tree =
    Node(key, 1, left, right)

  def internal(key: Int, level: Int, left: Tree, right: Tree): Tree =
    Node(key, level, left, right)

  def insert(t: Tree, x: Int): Tree = t match {
    case node: Node =>
      if (x < node.key) {
        node.left = insert(node.left, x)
      } else if (x > node.key) {
        node.right = insert(node.right, x)
      } else {
        // X == value(T) is unspecified
      }
      var node_ = skew(node)
      node_ = split(node_)
      node_
    case Nil =>
      Node(x, 1, Nil, Nil)
  }

  def delete(t: Tree, x: Int): Tree = t match {
    case node: Node =>
      if (x < node.key) {
        node.left = delete(node.left, x)
      } else if (x > node.key) {
        node.right = delete(node.right, x)
      } else {
        if (isLeaf(node)) {
          return node.right
        } else if (isNil(node.left)) {
          val temp = successor(node)
          node.right = delete(node.right, temp.key)
          node.key = temp.key
        } else {
          val temp = predecessor(node)
          node.left = delete(temp.left, temp.key)
          node.key = temp.key
        }
      }

      var node_ = decrease_level(node)
      node_ = skew(node_)
      node_.asInstanceOf[Node].right = skew(node_.asInstanceOf[Node].right)
      if (!isNil(node_.asInstanceOf[Node].right)) {
        node_.asInstanceOf[Node].right.asInstanceOf[Node].right = skew(
          node_.asInstanceOf[Node].right.asInstanceOf[Node].right
        )
      }
      node_ = split(node_)
      node_.asInstanceOf[Node].right = split(node_.asInstanceOf[Node].right)
      node_
    case Nil => t
  }

  private def decrease_level(t: Tree): Tree = t match {
    case node: Node =>
      val shouldBe = Math.min(level(node.left), level(node.right)) + 1
      if (shouldBe < node.level) {
        node.level = shouldBe
        if (shouldBe < level(node.right)) {
          node.right.asInstanceOf[Node].level = shouldBe
        }
      }
      node
    case _ => Nil
  }

  // Skew is a right rotation to replace a subtree containing a left horizontal link with one containing
  // a right horizontal link instead.
  //    L <-- T           L --> T
  //   / \     \   ===>  /     / \
  //  A   B     R       A     B   R
  private def skew(t: Tree): Tree = t match {
    case node: Node =>
      if (isNil(node.left)) node
      else if (level(node.left) == level(node)) {
        val left = node.left.asInstanceOf[Node]
        node.left = left.right
        left.right = node
        left
      } else {
        t
      }
    case _ => Nil
  }

  // Split is a left rotation and level increase to replace a subtree containing two or more consecutive
  // right horizontal links with one containing two fewer consecutive right horizontal links.
  //    T --> R --> X          R
  //   /     /         ===>   / \
  //  A     B                T   X
  //                        / \
  //                       A   B
  private def split(t: Tree): Tree = t match {
    case node: Node =>
      if (isNil(node.right) || isNil(node.right.asInstanceOf[Node].right)) node
      else if (level(node) == level(node.right.asInstanceOf[Node].right)) {
        val right = node.right.asInstanceOf[Node]
        node.right = right.left
        right.left = node
        right.level += 1
        right
      } else {
        t
      }
    case _ => Nil
  }

  private def predecessor(t: Tree): Node = {
    assert(!isLeaf(t), "calling `predecessor` on leaf")

    @tailrec
    def goRight(t: Tree): Node =
      t match {
        case node: Node if node.right == Nil => node
        case node: Node                      => goRight(node.right)
      }

    t match {
      case node: Node => goRight(node.left)
    }
  }

  private def successor(t: Tree): Node = {
    assert(!isLeaf(t), "calling `successor` on leaf")

    @tailrec
    def goLeft(t: Tree): Node =
      t match {
        case node: Node if node.left == Nil => node
        case node: Node                     => goLeft(node.left)
      }

    t match {
      case node: Node => goLeft(node.right)
    }
  }

  // level of every leaf node is one
  private def isLeaf(t: Tree): Boolean = t match {
    case node: Node => node.level == 1
    case _          => false
  }

  private def isNil(t: Tree): Boolean = t.isInstanceOf[Nil.type]

  private def level(t: Tree): Int = t match {
    case node: Node => node.level
    case _          => 0
  }
}

object Helper {
  private val colors =
    Array("#bea9de", "#87889c", "#546bab", "#2e4482", "#131862")

  case class Js(id: Int, parent: Option[Int], color: String) {
    override def toString: String = {
      val p = parent.map(_.toString).getOrElse("null")
      s"""{"id": $id, "parent": $p, "color": "$color"}""".stripMargin
    }
  }

  // level -> keys
  private def byLevels(t: Tree): Map[Int, List[(Int, Int)]] = {
    def inner(parent: Int, t: Tree, acc: Map[Int, List[(Int, Int)]]): Map[Int, List[(Int, Int)]] = t match {
      case node: Node =>
        val m    = acc.getOrElse(node.level, List.empty)
        var acc_ = acc.updated(node.level, m :+ (parent, node.key))
        acc_ = inner(node.key, node.left, acc_)
        acc_ = inner(node.key, node.right, acc_)
        acc_
      case Nil => acc
    }
    inner(0, t, Map.empty)
  }

  def stringify(t: Tree): String = {
    val map = byLevels(t)
    map.flatMap { case (level, nodes) =>
      nodes.map { case (parent, key) =>
        val p = if (parent == 0) None else Some(parent)
        Js(key, p, colors(level))
      }
    }.toList.mkString("[", ", ", "],")
  }
}

object Test extends App {
  // tree example from the paper https://user.it.uu.se/~arneande/ps/simp.pdf

  var root = {
    // level 1
    val n1  = Tree.leaf(1)
    val n3  = Tree.leaf(3)
    val n7  = Tree.leaf(7)
    val n5  = Tree.leaf(5, Nil, n7)
    val n9  = Tree.leaf(9)
    val n11 = Tree.leaf(11)
    val n13 = Tree.leaf(13)
    // level 2
    val n2  = Tree.internal(2, 2, n1, n3)
    val n8  = Tree.internal(8, 2, n5, n9)
    val n12 = Tree.internal(12, 2, n11, n13)
    // level 3
    val n10 = Tree.internal(10, 3, n8, n12)
    val n4  = Tree.internal(4, 3, n2, n10)

    n4
  }

  println(Helper.stringify(root))
  root = Tree.insert(root, 6)
  println(Helper.stringify(root))
  root = Tree.delete(root, 1)
  println(Helper.stringify(root))

  // Note: `stringify` produces JSON that can be used as input for the Treeviz JavaScript library, available at https://github.com/PierreCapo/treeviz, for visualization.
  //
  // [{"id": 4, "parent": null, "color": "#2e4482"}, {"id": 10, "parent": 4, "color": "#2e4482"}, {"id": 2, "parent": 4, "color": "#546bab"}, {"id": 8, "parent": 10, "color": "#546bab"}, {"id": 12, "parent": 10, "color": "#546bab"}, {"id": 1, "parent": 2, "color": "#87889c"}, {"id": 3, "parent": 2, "color": "#87889c"}, {"id": 5, "parent": 8, "color": "#87889c"}, {"id": 7, "parent": 5, "color": "#87889c"}, {"id": 9, "parent": 8, "color": "#87889c"}, {"id": 11, "parent": 12, "color": "#87889c"}, {"id": 13, "parent": 12, "color": "#87889c"}],
  // [{"id": 4, "parent": null, "color": "#2e4482"}, {"id": 10, "parent": 4, "color": "#2e4482"}, {"id": 2, "parent": 4, "color": "#546bab"}, {"id": 6, "parent": 10, "color": "#546bab"}, {"id": 8, "parent": 6, "color": "#546bab"}, {"id": 12, "parent": 10, "color": "#546bab"}, {"id": 1, "parent": 2, "color": "#87889c"}, {"id": 3, "parent": 2, "color": "#87889c"}, {"id": 5, "parent": 6, "color": "#87889c"}, {"id": 7, "parent": 8, "color": "#87889c"}, {"id": 9, "parent": 8, "color": "#87889c"}, {"id": 11, "parent": 12, "color": "#87889c"}, {"id": 13, "parent": 12, "color": "#87889c"}],
  // [{"id": 6, "parent": null, "color": "#2e4482"}, {"id": 10, "parent": 6, "color": "#2e4482"}, {"id": 4, "parent": 6, "color": "#546bab"}, {"id": 8, "parent": 10, "color": "#546bab"}, {"id": 12, "parent": 10, "color": "#546bab"}, {"id": 2, "parent": 4, "color": "#87889c"}, {"id": 3, "parent": 2, "color": "#87889c"}, {"id": 5, "parent": 4, "color": "#87889c"}, {"id": 7, "parent": 8, "color": "#87889c"}, {"id": 9, "parent": 8, "color": "#87889c"}, {"id": 11, "parent": 12, "color": "#87889c"}, {"id": 13, "parent": 12, "color": "#87889c"}],
}