数据结构之二叉平衡搜索树（AVL）

package datastruct.binarytree;

import java.util.Comparator;
import java.util.Objects;

public class AVLTree<E> {
    private static class Node<E>{
        E item;
        Node<E> parent;
        Node<E> left;
        Node<E> right;
        int height = 1;

        private Node(E item, Node<E> parent) {
            this.item = item;
            this.parent = parent;
        }

        private boolean isLeftChild(){
            return parent==null?false:parent.left==this;
        }

        private boolean isRightChild(){
            return parent==null?false:parent.right==this;
        }

        private boolean hasTwoChildren(){
            return left!=null&&right!=null;
        }

        private boolean isLeaf(){
            return left==null&&right==null;
        }

        private void setHeight(){
            int leftHeight = left==null?0:left.height;
            int rightHeight = right==null?0:right.height;
            this.height = (Math.max(leftHeight,rightHeight)+1);
        }

        private boolean isBalance(){
            int leftHeight = left==null?0:left.height;
            int rightHeight = right==null?0:right.height;
            return Math.abs(leftHeight-rightHeight)<=1;
        }

        private Node<E> tallerNode(){
            int leftHeight = left==null?0:left.height;
            int rightHeight = right==null?0:right.height;
            return (Math.max(leftHeight,rightHeight)==leftHeight)?left:right;
        }
    }

    private int size;

    private Node<E> root;

    private Comparator<E> comparator;

    public AVLTree() {
        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public int size(){
        return size;
    }

    public boolean isEmpty(){
        return size==0;
    }

    public boolean contains(E e){
        return node(e)!=null;
    }

    private Node<E> node(E e){
        Node<E> node = root;
        while (node!=null){
            int result = compare(e,node.item);
            if (result<0){
                node = node.left;
            }else if (result>0){
                node = node.right;
            }else {
                return node;
            }
        }
        return null;
    }

    private int compare(E e1, E e2) {
        return (comparator!=null)?comparator.compare(e1,e2):((Comparable<E>)e1).compareTo(e2);
    }

    public void clear(){
        root = null;
        size = 0;
    }

    public void add(E e){
        if (e==null) return;

        Node<E> node = root;
        if (node==null){
            root = new Node<>(e,null);
            size++;
            return;
        }

        int result = 0;
        Node<E> parent = null;
        while (node!=null){
            parent = node;
            result = compare(e,node.item);
            if (result<0){
                node = node.left;
            }else if (result>0){
                node = node.right;
            }else {
                node.item = e;
                return;
            }
        }

        Node<E> newNode = new Node<>(e,parent);
        if (result<0){
            parent.left = newNode;
        }else if (result>0){
            parent.right = newNode;
        }
        size++;

        afterAdd(newNode);
    }

    private void afterAdd(Node<E> node) {
        while ((node = node.parent)!=null){
            if (node.isBalance()){
                node.setHeight();
            }else {
                dealBalance(node);
                break;
            }
        }
    }

    private void dealBalance(Node<E> grand) {
        Node<E> parent = grand.tallerNode();
        Node<E> node = parent.tallerNode();

        if (parent.isLeftChild()){
            if (node.isLeftChild()){
                rightRotate(grand);
            }else {
                leftRotate(parent);
                rightRotate(grand);
            }
        }else {
            if (node.isRightChild()){
                leftRotate(grand);
            }else {
                rightRotate(parent);
                leftRotate(grand);
            }
        }
    }

    private void rightRotate(Node<E> grand) {
        Node<E> parent = grand.left;

        grand.left = parent.right;
        parent.right = grand;

        parent.parent = grand.parent;
        if (grand.isLeftChild()){
            grand.parent.left = parent;
        }else if (grand.isRightChild()){
            grand.parent.right = parent;
        }else {
            root = parent;
        }

        if (grand.left!=null){
            grand.left.parent = grand;
        }

        grand.parent = parent;

        grand.setHeight();
        parent.setHeight();

    }

    //error replace
    private void leftRotate(Node<E> grand) {
        Node<E> parent = grand.right;

        grand.right = parent.left;
        parent.left = grand;

        parent.parent = grand.parent;
        if (grand.isLeftChild()){
            grand.parent.left = parent;
        }else if (grand.isRightChild()){
            grand.parent.right = parent;
        }else {
            root = parent;
        }

        if (grand.right!=null){
            grand.right.parent = grand;
        }

        grand.parent = parent;

        grand.setHeight();
        parent.setHeight();
    }

    public void remove(E e){
        if (e==null) return;
        Node<E> node = node(e);
        if (node==null) return;

        if (node.hasTwoChildren()){
            Node<E> post = postNode(node);
            node.item = post.item;
            node = post;
        }

        if (node.isLeaf()){
            if (node.isLeftChild()){
                node.parent.left = null;
            }else if (node.isRightChild()){
                node.parent.right = null;
            }else {
                root = null;
            }
        }else {//error no deal parent replace.parent = node.parent;
            Node<E> replace = node.left!=null?node.left:node.right;
            replace.parent = node.parent;
            if (node.isLeftChild()){
                node.parent.left = replace;
            }else if (node.isRightChild()){
                node.parent.right = replace;
            }else {
                root = replace;
            }
//            node = replace;
        }
        size--;

        afterRemove(node);
    }

    private void afterRemove(Node<E> node) {
        while ((node = node.parent)!=null){
            if (node.isBalance()){
                node.setHeight();
            }else {
                dealBalance(node);
            }
        }
    }

    private Node<E> preNode(Node<E> node){
        Node<E> left = node.left;
        if (left!=null){
            while (left.right!=null){
                left = left.right;
            }
            return left;
        }

        Node<E> parent = node.parent;
        while (parent!=null&&parent.isLeftChild()){
            parent = parent.parent;
        }
        return parent;
    }

    private Node<E> postNode(Node<E> node){
        Node<E> right = node.right;
        if (right!=null){
            while (right.left!=null){
                right = right.left;
            }
            return right;
        }

        Node<E> parent = node.parent;
        while (parent!=null&&parent.isRightChild()){
            parent = parent.parent;
        }
        return parent;
    }

    public String toString(){
        StringBuilder stringBuilder = new StringBuilder();
        toString(root,stringBuilder,"");
        return stringBuilder.toString();
    }

    private StringBuilder toString(Node<E> node,StringBuilder stringBuilder,String fix){
        if (node==null) return null;
        toString(node.left,stringBuilder,fix+"L-->");
        stringBuilder.append(fix).append("["+node.item+"]").append("\n");
        toString(node.right,stringBuilder,fix+"R-->");
        return stringBuilder;
    }
}