

 This page uses content from Wikipedia. The original article was at AVL tree. The list of authors can be seen in the page history. As with Rosetta Code, the text of Wikipedia is available under the GNU FDL. (See links for details on variance)



In computer science, an AVL tree is a self-balancing binary search tree. In an AVL tree, the heights of the two child subtrees of any node differ by at most one; at no time do they differ by more than one because rebalancing is done ensure this is the case. Lookup, insertion, and deletion all take O(log n) time in both the average and worst cases, where n is the number of nodes in the tree prior to the operation. Insertions and deletions may require the tree to be rebalanced by one or more tree rotations. Note the tree of nodes comprise a set, so duplicate node keys are not allowed.

AVL trees are often compared with red-black trees because they support the same set of operations and because red-black trees also take O(log n) time for the basic operations. Because AVL trees are more rigidly balanced, they are faster than red-black trees for lookup-intensive applications. Similar to red-black trees, AVL trees are height-balanced, but in general not weight-balanced nor μ-balanced; that is, sibling nodes can have hugely differing numbers of descendants.



Task

Implement an AVL tree in the language of choice, and provide at least basic operations.




import scala.collection.mutable
 
class AVLTree[A](implicit val ordering: Ordering[A]) extends mutable.SortedSet[A] {
 
  if (ordering eq null) throw new NullPointerException("ordering must not be null")
 
  private var _root: AVLNode = _
  private var _size = 0
 
  override def size: Int = _size
 
  override def foreach[U](f: A => U): Unit = {
    val stack = mutable.Stack[AVLNode]()
    var current = root
    var done = false
 
    while (!done) {
      if (current != null) {
        stack.push(current)
        current = current.left
      } else if (stack.nonEmpty) {
        current = stack.pop()
        f.apply(current.key)
 
        current = current.right
      } else {
        done = true
      }
    }
  }
 
  def root: AVLNode = _root
 
  override def isEmpty: Boolean = root == null
 
  override def min[B >: A](implicit cmp: Ordering[B]): A = minNode().key
 
  def minNode(): AVLNode = {
    if (root == null) throw new UnsupportedOperationException("empty tree")
    var node = root
    while (node.left != null) node = node.left
    node
  }
 
  override def max[B >: A](implicit cmp: Ordering[B]): A = maxNode().key
 
  def maxNode(): AVLNode = {
    if (root == null) throw new UnsupportedOperationException("empty tree")
    var node = root
    while (node.right != null) node = node.right
    node
  }
 
  def next(node: AVLNode): Option[AVLNode] = {
    var successor = node
    if (successor != null) {
      if (successor.right != null) {
        successor = successor.right
        while (successor != null && successor.left != null) {
          successor = successor.left
        }
      } else {
        successor = node.parent
        var n = node
        while (successor != null && successor.right == n) {
          n = successor
          successor = successor.parent
        }
      }
    }
    Option(successor)
  }
 
  def prev(node: AVLNode): Option[AVLNode] = {
    var predecessor = node
    if (predecessor != null) {
      if (predecessor.left != null) {
        predecessor = predecessor.left
        while (predecessor != null && predecessor.right != null) {
          predecessor = predecessor.right
        }
      } else {
        predecessor = node.parent
        var n = node
        while (predecessor != null && predecessor.left == n) {
          n = predecessor
          predecessor = predecessor.parent
        }
      }
    }
    Option(predecessor)
  }
 
  override def rangeImpl(from: Option[A], until: Option[A]): mutable.SortedSet[A] = ???
 
  override def +=(key: A): AVLTree.this.type = {
    insert(key)
    this
  }
 
  def insert(key: A): AVLNode = {
    if (root == null) {
      _root = new AVLNode(key)
      _size += 1
      return root
    }
 
    var node = root
    var parent: AVLNode = null
    var cmp = 0
 
    while (node != null) {
      parent = node
      cmp = ordering.compare(key, node.key)
      if (cmp == 0) return node // duplicate
      node = node.matchNextChild(cmp)
    }
 
    val newNode = new AVLNode(key, parent)
    if (cmp <= 0) parent._left = newNode
    else parent._right = newNode
 
    while (parent != null) {
      cmp = ordering.compare(parent.key, key)
      if (cmp < 0) parent.balanceFactor -= 1
      else parent.balanceFactor += 1
 
      parent = parent.balanceFactor match {
        case -1 | 1 => parent.parent
        case x if x < -1 =>
          if (parent.right.balanceFactor == 1) rotateRight(parent.right)
          val newRoot = rotateLeft(parent)
          if (parent == root) _root = newRoot
          null
        case x if x > 1 =>
          if (parent.left.balanceFactor == -1) rotateLeft(parent.left)
          val newRoot = rotateRight(parent)
          if (parent == root) _root = newRoot
          null
        case _ => null
      }
    }
 
    _size += 1
    newNode
  }
 
  override def -=(key: A): AVLTree.this.type = {
    remove(key)
    this
  }
 
  override def remove(key: A): Boolean = {
    var node = findNode(key).orNull
    if (node == null) return false
 
    if (node.left != null) {
      var max = node.left
 
      while (max.left != null || max.right != null) {
        while (max.right != null) max = max.right
 
        node._key = max.key
        if (max.left != null) {
          node = max
          max = max.left
        }
      }
      node._key = max.key
      node = max
    }
 
    if (node.right != null) {
      var min = node.right
 
      while (min.left != null || min.right != null) {
        while (min.left != null) min = min.left
 
        node._key = min.key
        if (min.right != null) {
          node = min
          min = min.right
        }
      }
      node._key = min.key
      node = min
    }
 
    var current = node
    var parent = node.parent
    while (parent != null) {
      parent.balanceFactor += (if (parent.left == current) -1 else 1)
 
      current = parent.balanceFactor match {
        case x if x < -1 =>
          if (parent.right.balanceFactor == 1) rotateRight(parent.right)
          val newRoot = rotateLeft(parent)
          if (parent == root) _root = newRoot
          newRoot
        case x if x > 1 =>
          if (parent.left.balanceFactor == -1) rotateLeft(parent.left)
          val newRoot = rotateRight(parent)
          if (parent == root) _root = newRoot
          newRoot
        case _ => parent
      }
 
      parent = current.balanceFactor match {
        case -1 | 1 => null
        case _ => current.parent
      }
    }
 
    if (node.parent != null) {
      if (node.parent.left == node) {
        node.parent._left = null
      } else {
        node.parent._right = null
      }
    }
 
    if (node == root) _root = null
 
    _size -= 1
    true
  }
 
  def findNode(key: A): Option[AVLNode] = {
    var node = root
    while (node != null) {
      val cmp = ordering.compare(key, node.key)
      if (cmp == 0) return Some(node)
      node = node.matchNextChild(cmp)
    }
    None
  }
 
  private def rotateLeft(node: AVLNode): AVLNode = {
    val rightNode = node.right
    node._right = rightNode.left
    if (node.right != null) node.right._parent = node
 
    rightNode._parent = node.parent
    if (rightNode.parent != null) {
      if (rightNode.parent.left == node) {
        rightNode.parent._left = rightNode
      } else {
        rightNode.parent._right = rightNode
      }
    }
 
    node._parent = rightNode
    rightNode._left = node
 
    node.balanceFactor += 1
    if (rightNode.balanceFactor < 0) {
      node.balanceFactor -= rightNode.balanceFactor
    }
 
    rightNode.balanceFactor += 1
    if (node.balanceFactor > 0) {
      rightNode.balanceFactor += node.balanceFactor
    }
    rightNode
  }
 
  private def rotateRight(node: AVLNode): AVLNode = {
    val leftNode = node.left
    node._left = leftNode.right
    if (node.left != null) node.left._parent = node
 
    leftNode._parent = node.parent
    if (leftNode.parent != null) {
      if (leftNode.parent.left == node) {
        leftNode.parent._left = leftNode
      } else {
        leftNode.parent._right = leftNode
      }
    }
 
    node._parent = leftNode
    leftNode._right = node
 
    node.balanceFactor -= 1
    if (leftNode.balanceFactor > 0) {
      node.balanceFactor -= leftNode.balanceFactor
    }
 
    leftNode.balanceFactor -= 1
    if (node.balanceFactor < 0) {
      leftNode.balanceFactor += node.balanceFactor
    }
    leftNode
  }
 
  override def contains(elem: A): Boolean = findNode(elem).isDefined
 
  override def iterator: Iterator[A] = ???
 
  override def keysIteratorFrom(start: A): Iterator[A] = ???
 
  class AVLNode private[AVLTree](k: A, p: AVLNode = null) {
 
    private[AVLTree] var _key: A = k
    private[AVLTree] var _parent: AVLNode = p
    private[AVLTree] var _left: AVLNode = _
    private[AVLTree] var _right: AVLNode = _
    private[AVLTree] var balanceFactor: Int = 0
 
    def parent: AVLNode = _parent
 
    private[AVLTree] def selectNextChild(key: A): AVLNode = matchNextChild(ordering.compare(key, this.key))
 
    def key: A = _key
 
    private[AVLTree] def matchNextChild(cmp: Int): AVLNode = cmp match {
      case x if x < 0 => left
      case x if x > 0 => right
      case _ => null
    }
 
    def left: AVLNode = _left
 
    def right: AVLNode = _right
  }
 
}