AVL-Tree Implementation in Java

Question

I'm looking for feedback on this AVL tree implementation. Particularly on the overall structure, readability and abstraction. I've tried to keep all AVL specific operations inside the balancePath method to allow for easy conversion to a red/black tree with only minor changes to the rest of the code. Are there any further abstractions I could make to simplify the code even further?

public class AVLTreeDictionary<K extends Comparable<K>,V> implements Dictionary<K,V>
{
    /* Duplicate keys handling by overwriting value. Null keys
     * unsupported. It is assumed they will not be passed in.
     * Duplicate and null values supported.
     */
    private Node root;
    public AVLTreeDictionary()
    {
        root = null;
    }
    public V insert(K key, V value)
    {
        /* Adds a new entry to the tree. If the key already exists
         * overwrites the value and returns the old value. Performs
         * AVL rebalancing afterwords.
         */
        Path path = getPath(key);
        if(path.child() != null)
        {
            V previous = path.child().value;
            path.child().value = value;
            return previous;
        }
        path.setChild(new Node(key, value));
        path.pop();
        balancePath(path);
        return null;
    }
    public V delete(K key)
    {
        /* Deletes an entry from the tree if it exists.
         * If 0 children: Deletes entry.
         * If 1 child: Replaces entry with child.
         * If 2 children: Randomly selected between successor/predecessor
         * and copies values then deletes succ/pred and replaces it with
         * its one child if it exists. Performs AVL rebalancing afterwords.
         */
        Path path = getPath(key);
        Node node = path.child();
        if(node == null)
        {
            return null;
        }
        if(node.right == null)
        {
            path.setChild(node.left);
        }
        else if(node.left == null)
        {
            path.setChild(node.right);
        }
        else
        {
            boolean direction = Math.random() >= 0.5;
            path.push(direction);
            while(path.grandChild(!direction) != null)
            {
                path.push(!direction);
            }
            node.key = path.child().key;
            node.value = path.child().value;
            path.setChild(path.grandChild(direction));
        }
        path.pop();
        balancePath(path);
        return null;
    }
    public V search(K key)
    {
        // Standard binary search.
        Node search = root;
        while(search != null)
        {
            int compare = key.compareTo(search.key);
            if(compare == 0)
            {
                return search.value;
            }
            search = search.child(compare > 0);
        }
        return null;
    }
    private Path getPath(K key)
    {
        /* Returns a path from the root to the key specified or the
         * position it would be added in if it does not exist. Includes
         * the root edge.
         */
        Path path = new Path(new Edge(null, false));
        while(true)
        {
            if(path.child() == null)
            {
                return path;
            }
            int compare = key.compareTo(path.child().key);
            if(compare == 0)
            {
                return path;
            }
            path.push(compare > 0);
        }
    }
    private void rotate(Edge edge, boolean direction)
    {
        /* Rotates the child of the edge with its child
         * in the direction specified. It is assumed both the child
         * and grandchild are not null.
         */
        Edge rotate = new Edge(edge.child(), direction);
        edge.setChild(rotate.child());
        rotate.setChild(rotate.grandChild(!direction));
        edge.child().setChild(rotate.parent, !direction);
    }
    private int height(Node node)
    {
        if(node == null)
        {
            return -1;
        }
        else
        {
            return node.height;
        }
    }
    private void balancePath(Path path)
    {
        /* Follows a path up the tree performing AVL rebalancing
         * and height updates as needed.
         */
        while(path.peek() != null)
        {
            int previous = path.child().height;
            if(Math.abs(path.child().balance()) > 1)
            {
                boolean direction = path.child().balance() > 0;
                if(path.child().balance() * path.grandChild(direction).balance() < 0)
                {
                    path.push(direction);
                    rotate(path.peek(), !direction);
                    path.pop().grandChild(direction).updateHeight();
                }
                rotate(path.peek(), direction);
                path.grandChild(!direction).updateHeight();
            }
            if(path.pop().child().updateHeight() == previous)
            {
                break;
            }
        }
    }
    private class Node
    {
        /* This node class has get/set child functions that use a variable
         * to select which child. This allows us to avoid writing duplicate
         * code for the mirror image of operations. The balance of a node
         * refers to the difference in height between the right and left
         * subtrees. A positive balance is right-heavy, negative is left-heavy.
         * False = Left, True = Right.
         */
        public Node left;
        public Node right;
        public K key;
        public V value;
        public int height;
        public Node(K key, V value)
        {
            left = null;
            right = null;
            this.key = key;
            this.value = value;
            height = 0;
        }
        public Node child(boolean direction)
        {
            if(direction)
            {
                return right;
            }
            else
            {
                return left;
            }
        }
        public void setChild(Node child, boolean direction)
        {
            if(direction)
            {
                right = child;
            }
            else
            {
                left = child;
            }
        }
        public int balance()
        {
            return height(right) - height(left);
        }
        public int updateHeight()
        {
            height = Math.max(height(right), height(left)) + 1;
            return height;
        }
    }
    private class Edge
    {
        /* An edge is a connection between two nodes. It is represented by the parent
         * node and the directon to the child node. An edge may point to a null child.
         * There is also a special edge called the root edge which has the root as its
         * child, it is represented by a null parent. This representation allows us
         * to avoid writing duplicate code for handling operations involving the root.
         */
        public Node parent;
        public boolean direction;
        public Edge(Node parent, boolean direction)
        {
            this.parent = parent;
            this.direction = direction;
        }
        public Node child()
        {
            if(parent == null)
            {
                return root;
            }
            else
            {
                return parent.child(direction);
            }
        }
        public void setChild(Node child)
        {
            if(parent == null)
            {
                root = child;
            }
            else
            {
                parent.setChild(child, direction);
            }
        }
        public Node grandChild(boolean direction)
        {
            /* Returns the grandchild in the direction specified.
             * It is assumed this will not be called if the child
             * is null.
             */
            return child().child(direction);
        }
    }
    private class Path
    {
        /* A path is a list of adjacent edges going down the tree. All operations
         * are performed on the bottom-most edge of the path.
         */
        private Stack<Edge> stack;
        public Path(Edge edge)
        {
            stack = new LinkedListStack<Edge>();
            stack.push(edge);
        }
        public void push(boolean direction)
        {
            /* Extends the path downwards in the direction specified. It is
             * assumed this will not be called if the path points to a null
             * child.
             */
            stack.push(new Edge(child(), direction));
        }
        public Edge pop()
        {
            /* Retracts the bottom of the path upwards one edge and
             * returns the deleted edge.
             */
            return stack.pop();
        }
        public Edge peek()
        {
            return stack.peek();
        }
        public Node child()
        {
            return stack.peek().child();
        }
        public void setChild(Node child)
        {
            stack.peek().setChild(child);
        }
        public Node grandChild(boolean direction)
        {
            return stack.peek().grandChild(direction);
        }
    }
    public String toString()
    {
        return toString(root, "", false);
    }
    public String toString(Node node, String prefix, boolean direction)
    {
        String string = "";
        if(node != null)
        {
            string += toString(node.right, prefix + (direction ? "        " : "│       "), true);
            string += prefix + (direction ? "┌─────" : "└─────") + String.format("(%s,%s)\n", node.key, node.value);
            string += toString(node.left, prefix + (direction ? "│       " : "        "), false);
        }
        return string;
    }
}

Answer 1

• I'm going to start with what initially made me think the most, which is the fact that the height of a node is represented in two ways: Implicitly through the two subtrees of a node, and explicitly as an int field. Having interdependent state like this is always dangerous and makes code error prone, because a bug could cause the int field to contain a value different from the actual height of the node based on the structure of its subtrees. Replacing the field height with a method height() that recursively calculates the height based on the subtrees would make the code more robust and easier to hande. But of course, this would have a catastrophic impact on the performance, since updating the height values of a path would, if I'm not mistaken, be an \$O(\log_{2}n)\$ operation (with \$n\$ being the number of tree nodes, and \$\log_{2}n\$ therefore being the average depth), whereas recursively calculating the height of a node would require examining the complete subtrees, which would make it \$O(n)\$.

  So I think that, in this case, having interdependent state is justified. But for this, I would consider it all the more important to ensure that height never contains a wrong value by inseparably binding any modification of a node's height to a replacement of one of its child nodes, since the height of a node depends directly on the heights of its child nodes. Moreover, updating a node's height must also update the height of all of its ancestors as long as the height of a node changes due to the update.

  I see two ways to accomplish this. The most obvious way would be for each node to contain a reference not only to its children, but also to its parent, so that every time a node's child is replaced, it will update its own height and recursively cause its ancestors to update their height as well. The problem with this approach is that it would not only require more space, but also more time, because for each edge created or removed, an additional field has to be assigned. Furthermore, a rotation would then require three height update "waves", whereas one would suffice (starting from the new position of the rotated node). And since the whole point of this trouble is performance, solving the problem like this seems counterproductive.

  The other approach I can think of would be to define the methods for replacing a node's child not for a node, but only for an absolute path (i.e. making them methods of Path instead of Node). These methods would have to tamper directly with the fields Node.left and Node.right, but on the other hand, they can update the node's height and the height of all of its ancestors up to the root, if necessary (updateHeights could then also be a method of Path). And to avoid three update waves with rotations as mentioned in the previous paragraph, you could make rotate a bulk operation of Path that calls updateHeights only once after the rotation is complete, instead of making it a method of AVLTreeDictionary.

• A path as represented by your class Path can start anywhere. However, I could imagine that the code would be easier to handle if a path started at the root by definition. For one thing, this would make the change described in the previous point easier, and for another, it would solve the problem you yourself mentioned with having to write duplicate code for the root in a more elegant way than by introducing the concept of a "root edge", which to me seems more like a workaround than a solution, not least because the property direction of an edge is meaningless in the context of the root edge.

  The easiest way to implement this would probably to redesign the class Edge so that it stores not its parent, but its child. Then an empty path could be defined as pointing to the root. Your class Edge doesn't have any methods to access the parent node anyway, so maybe this would even improve overall readability.

  With an edge not having any reference to its parent, it will become impossible to make the parent node point to a different child using just the Edge. But this might actually not be a negative after all, because, taking into account what I've written in the previous point, changing a parent-child relationship without having the full path that leads from the root to these nodes would not be desirable anyway.

• In the class AVLTreeDictionary, the method search(K) duplicates most of the method getPath(K). Why not simply call getPath(K) from within search(K) and extract the value from the destination node of the returned path?

• Speaking of the method getPath(K), I would rewrite your while(true) loop to this:

  int compare;
  while (path.child() != null
          && (compare = key.compareTo(path.child().key)) != 0) {
      path.push(compare > 0);
  }
  return path;
  

  Because now, the termination of the loop depends on its termination condition, which I think makes the code less confusing than if the loop is terminated from within.

• I find the method names involving the word "child" in your class Path confusing. "Child" only makes sense in the context of a node (and possibly an edge), but with a path, I would find something with "destination" more fitting, for example nodeAtDestination(), or simply destination(), because with the word "child", I automatically think of the child of the node the path points to, and I always have to think twice when trying to understand something like somePath.child() or somePath.grandChild(someDirection).