Return to Revisions

I'm implementing some data structures for practice and would like some feedback on this AVL tree implementation. I've tried to keep all the AVL specific operations inside the balance path method so that it can be readily converted to a red/black tree with few changes. I've made some abstractions like the Edge and Path private classes which greatly cleaned up the code but I'm at a loss for how to clean up the balancePath method to make it more efficient and readable.

public class AVLTreeDictionary<K extends Comparable<K>,V> implements Dictionary<K,V>
{
 /* Duplicate keys handled by overwriting value.
 * Null keys unsupported, it is assumed they will not
 * be passed in. Null/duplicate 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 it overwrites the value and returns the old
 * value. Performs AVL rebalancing afterwards.
 */
 Path path = getPath(key);
 if(path.getChild() != null)
 {
 V previous = path.getChild().value;
 path.getChild().value = value;
 return previous;
 }
 path.pop().setChild(new Node(key, value));
 balancePath(path);
 return null;
 }
 public V delete(K key)
 {
 /* Deletes an entry from the tree if it exists and
 * returns the value.
 * - If 0 children deletes node.
 * - If 1 child replaces node with child.
 * - If 2 children randomly picks between successor/predecessor
 * and copies data to the node then deletes successor/predecessor
 * and replaces it with its 1 child if it exists.
 * If any nodes were deleted, performs AVL rebalancing of the tree.
 */
 Path path = getPath(key);
 Node node = path.getChild();
 if(node == null)
 {
 return null;
 }
 V value = node.value;
 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.getChild().getChild(!direction) != null)
 {
 path.push(!direction);
 }
 node.key = path.getChild().key;
 node.value = path.getChild().value;
 path.setChild(path.getChild().getChild(direction));
 }
 path.pop();
 balancePath(path);
 return value;
 }
 public V search(K key)
 {
 // Standard binary search for a key.
 Node search = root;
 while(search != null)
 {
 int compare = key.compareTo(search.key);
 if(compare == 0)
 {
 return search.value;
 }
 search = search.getChild(compare > 0);
 }
 return null;
 }
 private Path getPath(K key)
 {
 /* Returns a path from the root that has the key(or the position
 * it would be added in if it does not exist) as its child. Includes
 * the root edge.
 */
 Path path = new Path(new Edge(null, false));
 while(true)
 {
 if(path.getChild() == null)
 {
 return path;
 }
 int compare = key.compareTo(path.getChild().key);
 if(compare == 0)
 {
 return path;
 }
 path.push(compare > 0);
 }
 }
 private void rotate(Edge edge, boolean direction)
 {
 /* Rotates the child of the specified edge with its
 * grandchild in the direction specified. It is
 * assumed this function will only be called if the
 * child and grandchild are not null.
 */
 Edge rotate = new Edge(edge.getChild(), direction);
 edge.setChild(rotate.getChild());
 rotate.setChild(rotate.getChild().getChild(!direction));
 edge.getChild().setChild(rotate.parent, !direction);
 }
 private void balancePath(Path path)
 {
 while(path.peek() != null)
 {
 if(Math.abs(path.getChild().balanceFactor()) > 1)
 {
 boolean direction = path.getChild().balanceFactor() > 0;
 if(path.getChild().balanceFactor() * path.getChild().getChild(direction).balanceFactor() < 0)
 {
 path.push(direction);
 rotate(path.peek(), !direction);
 path.pop().getChild().getChild(direction).updateHeight();
 }
 rotate(path.peek(), direction);
 path.getChild().getChild(!direction).updateHeight();
 }
 path.pop().getChild().updateHeight();
 }
 }
 private class Node
 {
 /* The balance factor is the difference in height between
 * left and right subtrees. This class has get/set functions
 * for accessing child values using a vairbale to select which
 * child. This allows us to avoid writing duplicate code for
 * the mirror image of operations. 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 getChild(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 balanceFactor()
 {
 return (right != null ? right.height : -1) - (left != null ? left.height : -1);
 }
 public void updateHeight()
 {
 height = Math.max(left != null ? left.height : -1, right != null ? right.height : -1) + 1;
 }
 }
 private class Edge
 {
 /* An edge is a connection between two nodes. It is represented
 * by the parent node and the direction to the child node. An edge
 * may point to null children. There is a special edge called the root
 * edge that has the root node as its child, this is represented by
 * a null parent. This allows us to avoid writing special exceptions
 * for operations involving the root node.
 */
 public Node parent;
 public boolean direction;
 public Edge(Node parent, boolean direction)
 {
 this.parent = parent;
 this.direction = direction;
 }
 public Node getChild()
 {
 if(parent == null)
 {
 return root;
 }
 else
 {
 return parent.getChild(direction);
 }
 }
 public void setChild(Node child)
 {
 if(parent == null)
 {
 root = child;
 }
 else
 {
 parent.setChild(child, direction);
 }
 }
 }
 private class Path
 {
 /* A path is a set of adjacent edges going through the tree.
 * The child of a path refers to the child of the edge at the
 * bottom end of the path.
 */
 public Stack<Edge> stack;
 public Path(Edge edge)
 {
 stack = new LinkedListStack<Edge>();
 stack.push(edge);
 }
 public Node getChild()
 {
 return stack.peek().getChild();
 }
 public void setChild(Node child)
 {
 stack.peek().setChild(child);
 }
 public void push(boolean direction)
 {
 /* Extends the path downward in the direction of the parameter
 * passed in. It is assumed this function will not be called if
 * the path leads to a dead end.
 */
 stack.push(new Edge(stack.peek().getChild(), direction));
 }
 public Edge pop()
 {
 /* Retracts the path backward one edge and returns the deleted
 * edge.
 */
 return stack.pop();
 }
 public Edge peek()
 {
 return stack.peek();
 }
 }
 public String toString()
 {
 return toString(root, "", true);
 }
 public String toString(Node node, String prefix, boolean left)
 {
 String string = "";
 if(node != null)
 {
 string += toString(node.right, prefix + (left ? "│ " : " "), false);
 string += prefix + (left ? "└─────" : "┌─────") + String.format("(%s,%s)\n", node.key, node.value);
 string += toString(node.left, prefix + (left ? " " : "│ "), true);
 }
 return string;
 }
}