AVL tree rotation in Java

user1832681 picture user1832681 · Nov 30, 2012 · Viewed 25.5k times · Source

I want to implement the Java AVL tree and to rotate the tree left and right. I am not getting this.

Can anybody by looking at the code below tell me how can I possibly rotate the tree left and right and then use fix up with those two functions to balance the AVL tree?

I hope someone here can guide me through this.

import java.util.Random;
import java.util.SortedSet;
import java.util.TreeSet;

public class AVLTree<T> extends 
BinarySearchTree<AVLTree.Node<T>, T> implements SSet<T> {

    Random rand;

    public static class Node<T> extends BSTNode<Node<T>,T> {
        int h;  // the height of the node
    }

    public AVLTree() {
        sampleNode = new Node<T>();
        rand = new Random();
        c = new DefaultComparator<T>();
    }

    public int height(Node<T> u) {
        return (u == null) ? 0 : u.h;
    }

    public boolean add(T x) {
        Node<T> u = new Node<T>();
        u.x = x;
        if (super.add(u)) {
            for (Node<T> w = u; w != nil; w = w.parent) {
                // walk back up to the root adjusting heights
                w.h = Math.max(height(w.left), height(w.right)) + 1;
            }
        fixup(u);
        return true;
    }
    return false;
}

public void splice(Node<T> u) {
    Node<T> w = u.parent;
    super.splice(u);
    for (Node<T> z = u; z != nil; z = z.parent)
        z.h = Math.max(height(z.left), height(z.right)) + 1;            
    fixup(w);
}

public void checkHeights(Node<T> u) {
    if (u == nil) return;
    checkHeights(u.left);
    checkHeights(u.right);
    if (height(u) != 1 + Math.max(height(u.left), height(u.right))) 
        throw new RuntimeException("Check heights shows incorrect heights");
    int dif = height(u.left) - height(u.right);
    if (dif < -1 || dif > 1)
        throw new RuntimeException("Check heights found height difference of " + dif);
}

/**
 * TODO: finish writing this method
 * @param u
 */
public void fixup(Node<T> u) {
    while (u != nil) {
        int dif = height(u.left) - height(u.right);
        if (dif > 1) {
            // TODO: add code here to fix AVL condition 
            // on the path from u to the root, if  necessary
        } else if (dif < -1) {
            // TODO: add code here to fix AVL condition
            // on the path from u to the root, if necessary                 
        }
        u = u.parent;
    }
}

public Node rotateLeft() {
    return rotateLeft(u.parent);
}

public void rotateLeft(Node<T> u) {
    // TODO: Recompute height values at u and u.parent
}

public void rotateRight(Node<T> u) {
    // TODO: Recompute height values at u and u.parent
}

public static <T> T find(SortedSet<T> ss, T x) {
    SortedSet<T> ts = ss.tailSet(x);
    if (!ts.isEmpty()) {
        return ts.first(); 
    }
    return null;
}

/**
 * This just does some very basic correctness testing
 * @param args
 */
public static void main(String[] args) {
    AVLTree<Integer> t = new AVLTree<Integer>();
    Random r = new Random(0);
    System.out.print("Running AVL tests...");
    int n = 1000;
    for (int i = 0; i < n; i++) {
        t.add(r.nextInt(2*n));
        t.checkHeights(t.r);
    }
    for (int i = 0; i < n; i++) {
        t.remove(r.nextInt(2*n));
        t.checkHeights(t.r);
    }
    System.out.println("done");

    t.clear();
    System.out.print("Running correctness tests...");
    n = 100000;
    SortedSet<Integer> ss = new TreeSet<Integer>();
    Random rand = new Random();
    for (int i = 0; i < n; i++) {
        Integer x = rand.nextInt(2*n);
        boolean b1 = t.add(x);
        boolean b2 = ss.add(x);
        if (b1 != b2) {
            throw new RuntimeException("Adding " + x + " gives " + b2 
                    + " in SortedSet and " + b1 + " in AVL Tree");
        }
    }
    for (int i = 0; i < n; i++) {
        Integer x = rand.nextInt(2*n);
        Integer x1 = t.find(x);
        Integer x2 = find(ss, x);
        if (x1 != x2) {
            throw new RuntimeException("Searching " + x + " gives " + x2 
                    + " in SortedSet and " + x1 + " in AVL Tree");
        }
        ss.headSet(x);
    }
    for (int i = 0; i < n; i++) {
        Integer x = rand.nextInt(2*n);
        boolean b1 = t.remove(x);
        boolean b2 = ss.remove(x);
        if (b1 != b2) {
            throw new RuntimeException("Error (2): Removing " + x + " gives " + b2 
                    + " in SortedSet and " + b1 + " in AVL Tree");
        }
    }
    for (int i = 0; i < n; i++) {
        Integer x = rand.nextInt(2*n);
        Integer x1 = t.find(x);
        Integer x2 = find(ss, x);
        if (x1 != x2) {
            throw new RuntimeException("Error (3): Searching " + x + " gives " + x2 
                    + " in SortedSet and " + x1 + " in AVL Tree");
        }
        ss.headSet(x);
    }
    System.out.println("done");
 }
}

Answer

Trying picture Trying · Nov 10, 2013

Full AVL tree implementation:

public class AVLTree<T> {

    private AVLNode<T> root;

    private static class AVLNode<T> {

        private T t;
        private int height;
        private AVLNode<T> left;
        private AVLNode<T> right;

        private AVLNode(T t) {
            this.t = t;
            height = 1;
        }
    }

    public void insert(T value) {
        root = insert(root, value);
    }

    private AVLNode<T> insert(AVLNode<T> n, T v) {
        if (n == null) {
            n = new AVLNode<T>(v);
            return n;
        } else {
            int k = ((Comparable) n.t).compareTo(v);
            if (k > 0) {
                n.left = insert(n.left, v);
            } else {
                n.right = insert(n.right, v);
            }
            n.height = Math.max(height(n.left), height(n.right)) + 1;
            int heightDiff = heightDiff(n);
            if (heightDiff < -1) {
                if (heightDiff(n.right) > 0) {
                    n.right = rightRotate(n.right);
                    return leftRotate(n);
                } else {
                    return leftRotate(n);
                }
            } else if (heightDiff > 1) {
                if (heightDiff(n.left) < 0) {
                    n.left = leftRotate(n.left);
                    return rightRotate(n);
                } else {
                    return rightRotate(n);
                }
            } else;

        }
        return n;
    }

    private AVLNode<T> leftRotate(AVLNode<T> n) {
        AVLNode<T> r = n.right;
        n.right = r.left;
        r.left = n;
        n.height = Math.max(height(n.left), height(n.right)) + 1;
        r.height = Math.max(height(r.left), height(r.right)) + 1;
        return r;
    }

    private AVLNode<T> rightRotate(AVLNode<T> n) {
        AVLNode<T> r = n.left;
        n.left = r.right;
        r.right = n;
        n.height = Math.max(height(n.left), height(n.right)) + 1;
        r.height = Math.max(height(r.left), height(r.right)) + 1;
        return r;
    }

    private int heightDiff(AVLNode<T> a) {
        if (a == null) {
            return 0;
        }
        return height(a.left) - height(a.right);
    }

    private int height(AVLNode<T> a) {
        if (a == null) {
            return 0;
        }
        return a.height;
    }


}