package com.test.tree;

/**
 * 带有平衡条件的二叉查找树
 * */
public class AVLBinarySearchTree<T extends Comparable<? super T>> {
    
    /*内部类，定义二叉树中的节点结构*/
    private static class TreeNode<T>{
        private T data; //节点的值
        private TreeNode<T> lt; //节点左子树
        private TreeNode<T> rt; //节点右子树
        private int height; //用来记录节点的高度，进行单旋转或双旋转
        
        public TreeNode(T data) {
            this(data, null, null);
        }
        public TreeNode(T data, TreeNode<T> lt, TreeNode<T> rt) {
            this.data = data;
            this.lt = lt;
            this.rt = rt;
            this.height = 0;
        }
        public T getData() {
            return data;
        }
        public void setData(T data) {
            this.data = data;
        }
        public TreeNode<T> getLt() {
            return lt;
        }
        public void setLt(TreeNode<T> lt) {
            this.lt = lt;
        }
        public TreeNode<T> getRt() {
            return rt;
        }
        public void setRt(TreeNode<T> rt) {
            this.rt = rt;
        }
        public int getHeight() {
            return height;
        }
        public void setHeight(int height) {
            this.height = height;
        }
    }
    
    /**
     * 计算节点的高度
     * @param t 输入子树
     * @return 返回树的高度
     */
    public int height(TreeNode<T> t){
        return t==null ? -1 : t.height;
    }
    
    /**
     * 给树中添加节点
     * @param data 要插入的节点值
     * @param t 要插入的子树
     * @return 插入后形成的新的树
     */
    public TreeNode<T> insert(T data, TreeNode<T> t){
        if(t == null){
            //树为空
            return new TreeNode<T>(data, null ,null);
        }
        int compareReslt = compare(data, t.data);
        if(compareReslt < 0){
            //插入的数小于节点数，放入左子树中
            t.lt = insert(data, t.lt) ; //递归插入左子树
            
            //插入后检查当前节点的左右子树是否平衡
            if(height(t.lt)-height(t.rt) == 2){
                if(compare(data, t.lt.data) < 0){
                    t = rotateWithLeftChild(t); //单旋转
                }else if(compare(data, t.lt.data) > 0){
                    t = doubleWithLeftChild(t); //双旋转
                }
            }
        }else if(compareReslt > 0){
            //插入的数大于节点数，放入右子树中
            t.rt = insert(data, t.rt) ; //递归插入左子树
            
            //插入后检查当前节点的左右子树是否平衡
            if(height(t.rt)-height(t.lt) == 2){
                if(compare(data, t.rt.data) > 0){
                    t = rotateWithRightChild(t); //单旋转
                }else if(compare(data, t.rt.data) > 0){
                    t = doubleWithRightChild(t); //双旋转
                }
            }
        }
        t.height = Math.max(height(t.lt), height(t.rt)) + 1;
        return t;
    }
    private TreeNode<T> rotateWithRightChild(TreeNode<T> k2) {
        // TODO Auto-generated method stub
        TreeNode<T> k1 = k2.rt; 
        k2.rt = k1.lt; //左子树的右节点介于左子树根节点和根节点之间，赋值给根节点的左子树
        k1.lt = k2; //将根节点赋值给左节点的右节点
        k2.height = Math.max(height(k2.lt), height(k2.rt)) + 1;
        k1.height = Math.max(height(k1.rt), k2.height) + 1;
        return k1;
    }

    private TreeNode<T> doubleWithRightChild(TreeNode<T> k3) {
        // TODO Auto-generated method stub
        k3.rt = rotateWithLeftChild(k3.rt);
        return rotateWithRightChild(k3);
    }

    /**
     * 单旋转
     * @param t
     * @return
     */
    private TreeNode<T> rotateWithLeftChild(TreeNode<T> k2) {
        // TODO Auto-generated method stub
        TreeNode<T> k1 = k2.lt; //左子树的根节点赋给K1
        k2.lt = k1.rt; //左子树的右节点介于左子树根节点和根节点之间，赋值给根节点的左子树
        k1.rt = k2; //将根节点赋值给左节点的右节点
        k2.height = Math.max(height(k2.lt), height(k2.rt)) + 1;
        k1.height = Math.max(height(k1.lt), k2.height) + 1;
        return k1;
    }
    /**
     * 双旋转
     * @param t
     * @return
     */
    private TreeNode<T> doubleWithLeftChild(TreeNode<T> k3) {
        // TODO Auto-generated method stub
        k3.lt = rotateWithRightChild(k3.lt);
        return rotateWithLeftChild(k3);
    }

    /**
     * 比较两个值是否相等
     * @param data1
     * @param data2
     * @return
     */
    public int compare(T data1, T data2){
        return data1.compareTo(data2);
    }
    
    /*中序遍历*/
    public void printTree(TreeNode<T> t){
        if(t != null){
            printTree(t.lt);
            System.out.print(t.data+"、");
            printTree(t.rt);
        }
    }
    
    public static void main(String[] args) {
        AVLBinarySearchTree<Integer> aVLBinarySearchTree = new AVLBinarySearchTree<Integer>();
        TreeNode<Integer> node = new TreeNode<Integer>(8);
        node = aVLBinarySearchTree.insert(6, node);
        node = aVLBinarySearchTree.insert(16, node);
        node = aVLBinarySearchTree.insert(13, node);
        node = aVLBinarySearchTree.insert(19, node);
        node = aVLBinarySearchTree.insert(7, node);
        node = aVLBinarySearchTree.insert(21, node);
        node = aVLBinarySearchTree.insert(23, node);
        aVLBinarySearchTree.printTree(node.lt);
        aVLBinarySearchTree.printTree(node.rt);
    }
}