平衡二叉树Java实现

平衡二叉树Java实现

static class TreeNode{

        public int val;
        public TreeNode left;
        public TreeNode right;

        public TreeNode(int val){
            this.val = val;
        }

    }

    static class BTree{

        private TreeNode root ;
        private int length=0;
        //添加子节点
        public void add(int val){
            if (root==null){
                root = new TreeNode(val);
            }else {
                TreeNode bTree = root;
                TreeNode parent = root;
                //寻找位置插入
                while (bTree!=null){
                    parent = bTree;
                    if (bTree.val < val){
                        bTree = bTree.right;
                    }else {
                        bTree = bTree.left;
                    }
                }
                if (parent.val < val){
                    parent.right = new TreeNode(val);
                }else {
                    parent.left = new TreeNode(val);
                }

                //插值看平衡是否被破坏，若被破坏则需要旋转二叉树保存平衡
                int count=1;
                while (!isBTree(root)){
                    System.out.println("正在进行第"+(count++)+"次旋转" );
                    //存储父节点和祖父节点方便旋转
                    TreeNode[] tns = new TreeNode[2];
                    TreeNode treeNode = root;
                    //找到最小的不平衡树 tns[1]
                    while (!isBTree(treeNode)){
                        tns[0] = tns[1];
                        tns[1] = treeNode;
                        if (!isBTree(treeNode.left)){
                            treeNode = treeNode.left;
                        }else if(!isBTree(treeNode.right)) {
                            treeNode = treeNode.right;
                        }else {
                            break;
                        }
                    }
                    //通过对tns[1]子树的高度 和 tns[1]的子节点来判断旋转类型
                    int hl = height(tns[1].left);
                    int hr = height(tns[1].right);
                    int type = 0; // 0:ll型,1:rr型,2:rl型,3:lr型
                    if (hl>hr && tns[1].left.right!=null){
                        type = 3;
                    }
                    if (hl<hr && tns[1].right.left==null){
                        type = 1;
                    }
                    if(hl<hr && tns[1].right.left!=null){
                        type = 2;
                    }
                    //旋转
                    rotate(type,tns[1],tns[0]);
                }

            }
        }

        /**
         * 旋转函数
         * @param type // 0ll,1rr,2rl,3lr
         * @param parent 父节点
         * @param parentsPatent 祖父节点
         */
        private void rotate(int type,TreeNode parent,TreeNode parentsPatent){
            TreeNode treeNode ;
            switch (type){
                case 0:  //ll型
                    treeNode = parent.left;
                    treeNode.right = parent;
                    parent.left = null;
                    break;
                case 1: //rr
                    treeNode = parent.right;
                    treeNode.left = parent;
                    parent.right = null;
                    break;
                case 2: // rl
                    treeNode = parent.right.left;
                    TreeNode r = parent.right;
                    parent.right = treeNode.left;
                    r.left = treeNode.right;
                    treeNode.left = parent;
                    treeNode.right = r;
                    break;
                case 3: //lr
                    treeNode = parent.left.right;
                    TreeNode l = parent.left;
                    parent.left = treeNode.right;
                    l.right = treeNode.left;
                    treeNode.right = parent;
                    treeNode.left = l;
                    break;
                default:
                    throw new RuntimeException("旋转错误");
            }
            //若果父节点是根节点，则需要更新root
            if (parent==root){
                root = treeNode;
            }else {
                if(parentsPatent.left == parent){
                    parentsPatent.left = treeNode;
                }else {
                    parentsPatent.right = treeNode;
                }
            }
        }

        public boolean isBTree(){
            return isBTree(root);
        }
        public int height(){
            return height(root);
        }
        //判断是否是平衡二叉树
        private boolean isBTree(TreeNode treeNode){
            if(treeNode == null){
                return true;
            }
            return isBTree(treeNode.left) && isBTree(treeNode.right) && Math.abs(height(treeNode.left)-height(treeNode.right))<=1;
        }
        //计算二叉树的高度
        private int height(TreeNode treeNode){
            if(treeNode==null){
                return 0;
            }
            return Math.max(height(treeNode.left),height(treeNode.right))+1;
        }
        //计算该二叉树的元素个数
        public int getLength(){
            length = 0;
            calLength(root);
            return length;
        }
        //计算该二叉树的元素个数
        private void calLength(TreeNode root){
            if (root!=null){
                length++;
                calLength(root.left);
                calLength(root.right);
            }
        }
        //将二叉树先序遍历转换成数组
        public int[] bTreeToArray(){
            int[] rs = new int[getLength()];
            length = 0;
            preorderTraversal( root,rs);
            return rs;
        }

        //先序遍历二叉树
        private void preorderTraversal(TreeNode root,int[] rs){
            if (root != null){
                preorderTraversal(root.left,rs);
                rs[length++] = root.val;
                preorderTraversal(root.right,rs);
            }
        }

        public void preorderTraversal(){
            System.out.println("-----------------------------------------------------------------------");
            preorderTraversal( root);
            System.out.println("-----------------------------------------------------------------------");
        }
        private void preorderTraversal(TreeNode root){
            if (root != null){
                preorderTraversal(root.left);
                System.out.print(root.val+"  ");
                preorderTraversal(root.right);
            }
        }

    }

    //测试验证函数
    public static void main(String[] args) {
        BTree bTree ;
        int nz=100000;
        int ng=50;
        int[] flag;
        for (int i = 0; i < nz; i++) {
            flag = new int[101];
            bTree = new BTree();
            Random random = new Random(System.currentTimeMillis());
            int[] sources = new int[ng];
            for (int j = 0; j < ng; j++) {
                sources[j] = random.nextInt(100);
                flag[sources[j]] +=1 ;
                System.out.println(sources[j]);
                try {
                    bTree.add(sources[j]);
                }catch (StackOverflowError e){
                    System.out.println(Arrays.toString(bTree.bTreeToArray()));
                    return;
                }
            }
            if(!bTree.isBTree()){
                throw new RuntimeException("不是平衡二叉树,算法有误");
            }
            System.out.println(Arrays.toString(sources));
            int[] ints = bTree.bTreeToArray();
            if (bTree.getLength()!=ng){
                System.out.println("bTreeLength ="+bTree.getLength() +",ng="+ng);
                bTree.preorderTraversal();
                throw new RuntimeException("二叉搜索树异常");
            }else {
                //平衡二叉树数组
                System.out.println(Arrays.toString(ints));
                for (int j = 0; j < ng-1; j++) {
                    flag[ints[j]] --;
                    if (ints[j]> ints[j+1]){
                        throw new RuntimeException("二叉排序树数组生成异常");
                    }
                    if(flag[ints[j]] < 0){
                        throw new RuntimeException("二叉排序树插入异常");
                    }
                }
                if(flag[ints[ng-1]] <= 0){
                    throw new RuntimeException("二叉排序树插入异常");
                }
            }

        }
    }