二叉平衡树(AVL tree)的java实现

二叉平衡树(AVL tree)的java实现

package com.jiading;

/*
平衡二叉树(AVL tree)的java实现
AVL树的特点：
平衡二叉树，是一种二叉排序树，其中每个结点的左子树和右子树的高度差至多等于1。它是一种高度平衡的二叉排序树。
 */
public class AVLTree {
    private Node root;

    /*
    定义Node内部类
     */
    private static class Node {
        private int key;
        //balance表示左右子树的高度差
        private int balance;
        private int height;
        private Node left;
        private Node right;
        //注意，AVL树的节点要有指向parent的指针
        private Node parent;

        public Node(int key, Node parent) {
            this.key = key;
            this.parent = parent;
        }
    }

    /*
    Node类定义完成
    insert,返回值表示插入是否成功.不能重复，重复就插入失败
     */
    public boolean insert(int key) {
        if (root == null) {
            root = new Node(key, null);
            return true;
        }
        //n表示待插入位置，现在在找位置
        Node n = root;
        while(true){
            if(n.key==key){
                return false;
            }
            Node parent=n;
            boolean goLeft=n.key>key;
            n=goLeft?n.left:n.right;
            //n==null就表示可以插入到这里了
            if(n==null){
                if(goLeft){
                    parent.left=new Node(key,parent);
                }else{
                    parent.right=new Node(key,parent);
                }
                rebalance(parent);
                break;
            }
        }
        return true;
    }
    private void delete(Node node){
        if(node.left==null&&node.right==null){
            if(node.parent==null){
                //为什么这里写root,因为这种情况只可能是root才有
                root=null;
            }else{
                Node parent=node.parent;
                if(parent.left==node){
                    parent.left=null;
                }else{
                    parent.right=null;
                }
                rebalance(parent);
            }
            return;
        }
        /*
        非叶子节点的删除不会导致AVL树在删除的位置失衡，因为又从叶子节点补上了
        需要调整的在叶子节点的位置，会由delete(child)调整的
         */
        if(node.left!=null){
            Node child=node.left;
            /*
            找左孩子的最右叶子节点，把它的值换上来，然后删除原来的叶子节点
             */
            while(child.right!=null){
                child=child.right;
            }
            node.key=child.key;
            delete(child);
        }else{
            Node child=node.right;
            while(child.left!=null){
                child=child.left;
            }
            node.key=child.key;
            delete(child);
        }
    }
    public  void delete(int delKey){
        if(root==null){
            return;
        }
        Node child=root;
        while(child!=null){
            Node node=child;
            child=delKey>node.key?node.right:node.left;
            if(delKey==node.key){
                delete(node);
                return;
            }
        }
    }
    /*
    计算节点高度
     */
    private int height(Node n){
        if(n==null){
            return -1;
        }
        return n.height;
    }
    /*
    重新计算节点的高度,子节点高度最高的加一
     */
    private void reheight(Node node){
        if(node!=null){
            node.height=1+Math.max(height(node.left),height(node.right));
        }
    }
    /*
    ...是可变参数的意思，和Node[]等效
     */
    private void setBalance(Node... nodes){
        for(Node n:nodes){
            reheight(n);
            n.balance=height(n.right)-height(n.left);
        }

    }
    /*
    将子树往左转
     */
    private Node rotateLeft(Node a){
        Node b=a.right;
        b.parent=a.parent;
        a.right=b.left;
        if(a.right!=null){
            a.right.parent=a;
        }
        b.left=a;
        a.parent=b;
        /*
        对parent节点进行设置
         */
        if(b.parent!=null){
            if(b.parent.right==a){
                b.parent.right=b;
            }else{
                b.parent.left=b;
            }
        }
        setBalance(a,b);
        return b;
    }
    private Node rotateRight(Node a){
        Node b=a.left;
        b.parent=a.parent;
        a.left=b.right;
        if(a.left!=null){
            a.left.parent=a;
        }
        b.right=a;
        a.parent=b;
        /*
        对parent节点进行设置
         */
        if(b.parent!=null){
            if(b.parent.right==a){
                b.parent.right=b;
            }else{
                b.parent.left=b;
            }
        }
        setBalance(a,b);
        return b;
    }
    private Node rotateLeftThenRight(Node n){
        //先旋转左子树
        n.left=rotateLeft(n.left);
        //然后旋转当前节点
        return rotateRight(n);
    }
    private Node rotateRightThenLeft(Node n){
        n.right=rotateRight(n.right);
        return rotateLeft(n);
    }
    private void rebalance(Node n){
        setBalance(n);
        if(n.balance==-2){
            if(height(n.left.left)>=height(n.left.right)){
                n=rotateRight(n);
            }else{
                n=rotateLeftThenRight(n);
            }
        }else if(n.balance==2){
            if(height(n.right.right)>=height(n.right.left)){
                n=rotateLeft(n);
            }else{
                n=rotateRightThenLeft(n);
            }
        }
        /*
        自上而下地调整
         */
        if(n.parent!=null){
            rebalance(n.parent);
        }else{
            root=n;
        }
    }
    public void printBalance(){
        printBalance(root);
    }
    public void printBalance(Node n){
        if(n!=null){
            printBalance(n.left);
            System.out.printf("%s",n.balance);
            printBalance(n.right);
        }
    }

    public static void main(String[] args) {
        AVLTree tree=new AVLTree();
        System.out.println("Inserting values 1 to 10");
        for(int i=1;i<10;i++){
            tree.insert(i);
        }
        System.out.println("Printing balance");
        tree.printBalance();
    }
}