package dataStructures;
public class AVLTree extends BSTree {
private static class AVLNode extends Node {
private int height;
AVLNode (Object e, Node left, Node right) {
super(e, left, right);
setHeight();
}
void setHeight() {
height = 1+Math.max(height(left),height(right));
}
int height(Node n) {
return (n == null ? -1 : ((AVLNode) n).height);
}
int balanceValue() {
return height(right) - height(left);
}
}
Node add(Node r, Object e) {
if (r == null) {
r = new AVLNode(e, null, null);
++size;
} else {
r = super.add(r,e);
r = rebalance(r);
}
return r;
}
Node remove(Node r, Object e) {
r = super.remove(r,e);
r = rebalance(r);
return r;
}
private Node rebalance(Node r) {
if (r == null) return r;
int balance = ((AVLNode)r).balanceValue();
if (balance == -2) {
if (((AVLNode)r.left).balanceValue() == 1)
r.left = rotateRightChild(r.left);
r = rotateLeftChild(r);
} else if (balance == 2) {
if (((AVLNode)r.right).balanceValue() == -1)
r.right = rotateLeftChild(r.right);
r = rotateRightChild(r);
}
((AVLNode)r).setHeight();
return r;
}
Node rotateLeftChild(Node r) {
r = super.rotateLeftChild(r);
((AVLNode) r.right).setHeight();
((AVLNode) r).setHeight();
return r;
}
Node rotateRightChild(Node r) {
r = super.rotateRightChild(r);
((AVLNode) r.left).setHeight();
((AVLNode) r).setHeight();
return r;
}
}