AVLuto

package adt.avltree;

import java.util.Arrays;

import adt.bst.BSTImpl;
import adt.bst.BSTNode;
import adt.bt.Util;

/**
 *
 * Performs consistency validations within a AVL Tree instance
 *
 * @author Claudio Campelo
 *
 * @param <T>
 */
public class AVLTreeImpl<T extends Comparable<T>> extends BSTImpl<T> implements
        AVLTree<T> {

    // TODO Do not forget: you must override the methods insert and remove
    // conveniently.

    // AUXILIARY
    protected int calculateBalance(BSTNode<T> node) {
        return height((BSTNode<T>)node.getLeft())-height((BSTNode) node.getRight());
    }

    protected void toLeft(BSTNode<T> node) {
        BSTNode<T> center = Util.leftRotation(node);
        if(node.equals(this.root)) this.root = center;
    }

    protected void toRight(BSTNode<T> node) {
        BSTNode<T> center = Util.rightRotation(node);
        if(node.equals(this.root)) this.root = center;
    }
    // AUXILIARY
    protected void rebalance(BSTNode<T> node) {
        if(node != null) {
            int balance = this.calculateBalance(node);

            if(balance > 1) {
                int childBalance = calculateBalance((BSTNode<T>)node.getLeft());
                if(childBalance > 0) {
                    toRight(node);
                }else {
                    toLeft((BSTNode<T>)node.getLeft());
                    toRight(node);
                }
            }
            if(balance < -1) {
                int childBalance = calculateBalance((BSTNode<T>)node.getRight());
                if(childBalance < 0) {
                    toLeft(node);
                }else {
                    toRight((BSTNode<T>)node.getRight());
                    toLeft(node);
                }
            }

        }

    }

    // AUXILIARY
    protected void rebalanceUp(BSTNode<T> node) {
        if(node != null) {
            rebalance(node);
            rebalanceUp((BSTNode<T>)node.getParent());
        }
    }

    @Override
    public void insert(T element) {
        super.insert(element);
        BSTNode<T> node = search(element);
        rebalanceUp(node);
    }

    @Override
    public void remove(T element) {
       if (element != null) {
          BSTNode<T> no = search(element);
          remove(no);
       }
    }

    private void remove(BSTNode<T> no) {
       if(!no.isEmpty() && !no.equals(new BSTNode<T>())) {
          if(no.isLeaf()) {
              super.replace(no, new BSTNode());
              rebalanceUp((BSTNode<T>)no.getParent());
          }else if(no.getRight().isEmpty()) {
              super.replace(no,(BSTNode<T>) no.getLeft());
              this.rebalanceUp((BSTNode<T>)no.getParent());
          }else if(no.getLeft().isEmpty()) {
              super.replace(no, (BSTNode<T>) no.getRight());
              rebalanceUp((BSTNode<T>)no.getParent());
          }else {
              BSTNode<T> sucess = sucessor(no.getData());
              no.setData(sucess.getData());
              remove(sucess);
          }
       }
    }
}