java实现AVL树

AVL树是平衡性要求非常高的二叉查找树，查找效率很高，也很复杂。

写完之后终于明白java的hashmap为什么用红黑树而不用AVL树了。

public class BinaryAVLTree<K extends Comparable<K>> {

    private AVLNode<K> root;
    private int size;

    public AVLNode<K> getRoot() {
        return root;
    }

    public void setRoot(AVLNode<K> root) {
        this.root = root;
    }

    /**
     * 初步认为：增删节点，需要更新height的节点太多，选择删掉height这个属性
     * 在需要判断平衡的时候递归计算，可以降低程序的复杂度
     *
     * @param <K>
     */
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    static class AVLNode<K> {
        private K key;
        private AVLNode<K> left;
        private AVLNode<K> right;
        private AVLNode<K> parent;
//        private int height;
    }

    /**
     * 插入的非递归实现
     * 单纯的插入不难，问题是插入后如果树失去平衡，需要做一次自平衡
     * 插入一个节点，被插入节点的父节点不会失去平衡，失去平衡的可能是祖父节点
     * 所以要从小往上回溯，找到最小的非平衡树进行自平衡
     *
     * @param key
     */
    public void insert(K key) {
        if (root == null) {
            root = new AVLNode<>(key, null, null, null);
            System.out.println(root);
            size++;
            return;
        }
        AVLNode<K> cursor = root;
        //被更新节点的所有递归父节点的平衡都有可能被打破
        while (true) {
            if (key.compareTo(cursor.key) < 0) {
                if (cursor.left == null) {
                    cursor.left = new AVLNode<>(key, null, null, cursor);
                    size++;
                    break;
                }
                cursor = cursor.left;
            } else if (key.compareTo(cursor.key) > 0) {
                if (cursor.right == null) {
                    cursor.right = new AVLNode<>(key, null, null, cursor);
                    size++;
                    break;
                }
                cursor = cursor.right;
            } else {
                //待插入数据和已有数据相同，如果有value则更新，没有则可以跳出循环
                break;
            }
        }

        //从下到上遍历，判断是否需要做自平衡
        AVLNode<K> cp = cursor.parent;
        while (cp != null) {
            if (getBalanceFactor(cp) > 1 || getBalanceFactor(cp) < -1) {
                selfBalance(cp);
                break;
            }
            cp = cp.parent;
        }


    }
    //todo

    /**
     * 删除跟二叉查找树类似，只不过删完需要做自平衡
     *
     * @param key
     */
    public void remove(K key) {
        AVLNode<K> cursor = this.root;
        if (cursor == null) {
            return;
        }
        while (cursor != null) {
            K k = cursor.key;
            if (key.compareTo(k) < 0) {
                cursor = cursor.left;
            } else if (key.compareTo(k) > 0) {
                cursor = cursor.right;
            } else {
                doRemove(cursor);
                return;
            }
        }
    }

    /**
     * 删完需要做自平衡
     *
     * @param cursor
     */
    private void doRemove(AVLNode<K> cursor) {
        //cursor是一个叶子节点
        if (cursor.left == null && cursor.right == null) {
            //截断了父节点的左右孩子指针，但没有取消cursor的父指针，因为后面还需要用到
            if (cursor == cursor.parent.left) {
                cursor.parent.left = null;
            } else {
                cursor.parent.right = null;
            }
        } else if (cursor.right == null) {
            if (cursor == cursor.parent.left) {
                cursor.parent.left = cursor.left;
            } else {
                cursor.parent.right = cursor.left;
            }
            //不管是左孩子还是右孩子，孩子的父亲是没问题的
            cursor.left.parent=cursor.parent;
        } else if (cursor.left == null) {
            if (cursor == cursor.parent.left) {
                cursor.parent.left = cursor.right;
            } else {
                cursor.parent.right = cursor.right;
            }
            cursor.right.parent=cursor.parent;
        } else {
            //左右都不为null
            //选择左子树最大节点，放在自己的位置
            //左子树最大节点删除
            AVLNode<K> maxLeft = findMax(cursor.left);
            remove(maxLeft.key);
            cursor.setKey(maxLeft.key);
        }

        //从下到上遍历，判断是否需要做自平衡
        AVLNode<K> cp = cursor.parent;
        while (cp != null) {
            if (getBalanceFactor(cp) > 1 || getBalanceFactor(cp) < -1) {
                selfBalance(cp);
                break;
            }
            cp = cp.parent;
        }
    }

    public AVLNode<K> findMax(AVLNode<K> root) {
        if (root == null) {
            return null;
        }
        AVLNode<K> cursor = root;
        while (true) {
            if (cursor.right == null) {
                return cursor;
            } else {
                cursor = cursor.right;
            }
        }
    }


    public AVLNode<K> selfBalance(AVLNode<K> node) {
        //左-右大于1，左高
        if (getBalanceFactor(node) > 1) {
            //接下来判断左子树的哪个孩子高
            int leftFactor = getBalanceFactor(node.left);
            //左左高，LL
            if (leftFactor > 0) {
                return rightRotate(node);
                //左右高：LR
            } else {
                return leftRightRotate(node);
            }
            //左-右小于-1，右高
        } else if (getBalanceFactor(node) < -1) {
            int rightFactor = getBalanceFactor(node.right);
            //右右高：RR
            if (rightFactor < 0) {
                return leftRotate(node);
                //右左：RL
            } else {
                return rightLeftRotate(node);
            }
        } else {
            //当前是平衡的
            return node;
        }
    }

    /**
     * 旋转的是失去平衡的最小子树
     * 物极必反，RR左旋，当前节点右子树的右侧超高了
     * 共涉及3个节点，3个指针：
     * 1.最小非平衡子树根节点root,
     * 2.root的右子节点rightChild,
     * 3.rightChild的左孩子rightChild.left
     * 从下到上依次进行处理
     * 1.
     */
    public AVLNode<K> leftRotate(AVLNode<K> root) {
        System.out.println("leftRotate");
        //这个未来是新根
        AVLNode<K> rightChild = root.right;
        //1.rightChild的左孩子,如果右孩子有左子节点，将它送给原根做右节点，一并确认父子关系
        root.right = rightChild.left;
        //更新右孩子的左孩子的父指针
        if (rightChild.left != null) {
            rightChild.left.parent = root;
        }
        //2.root的右子节点rightChild，可能设计到树的根节点root
        //我是父亲的孩子，父亲的孩子是我
        rightChild.parent = root.parent;
        //如果root是整棵树的根节点，需要更新树的root
        if (root.parent == null) {
            this.root = rightChild;
            //如果root是原父亲的左孩子，把这个位置让给新孩子rightChild
        } else if (root == root.parent.left) {
            root.parent.left = rightChild;
        } else {
            root.parent.right = rightChild;
        }
        //3.root处理：原来的根成为新根的左节点
        //新根成为老root的父亲，父子关系一并确定
        rightChild.left = root;
        //更新父指针
        root.parent = rightChild;

        return rightChild;
    }

    /**
     * LL右旋:左节点的左孩子超高了
     * 操作与左旋正好相反
     */
    public AVLNode<K> rightRotate(AVLNode<K> root) {
        System.out.println("rightRotate");
        //新根
        AVLNode<K> leftChild = root.left;
        //如果左孩子有右子树，则将它送给原根节点做左孩子
        root.left = leftChild.right;
        if (leftChild.right != null) {
            leftChild.right.parent = root;
        }
        leftChild.parent = root.parent;
        if (root.parent == null) {
            this.root = leftChild;
        } else if (root.parent.left == root) {
            root.parent.left = leftChild;
        } else {
            root.parent.right = leftChild;
        }

        leftChild.right = root;
        root.parent = leftChild;

        return leftChild;
    }

    /**
     * LR:左右旋，左孩子的右子树超高了，需要先对左孩子左旋，然后自己右旋
     *
     * @param root
     * @return
     */
    public AVLNode<K> leftRightRotate(AVLNode<K> root) {
        System.out.println("leftRightRotate");
        AVLNode<K> leftChild = root.left;
        AVLNode<K> newLeft = leftRotate(leftChild);
        root.left = newLeft;
        newLeft.parent = root;
        //新根
        return rightRotate(root);
    }

    /**
     * RL:右左旋，右孩子的左子树超高了，需要先对右孩子右旋，然后自己左旋
     */

    public AVLNode<K> rightLeftRotate(AVLNode<K> root) {
        System.out.println("rightLeftRotate");
        AVLNode<K> rightChild = root.right;
        AVLNode<K> newRight = rightRotate(rightChild);
        root.right = newRight;
        newRight.parent = root;
        //新根
        return leftRotate(root);
    }

    /**
     * 获取某个节点的平衡因子
     * 平衡因子： 某个结点的左子树的高度减去右子树的高度得到的差值。
     *
     * @param root
     * @return
     */
    public int getBalanceFactor(AVLNode<K> root) {

        int leftHeight = -1;
        int rightHeight = -1;
        if (root.left != null) {
            leftHeight = getHeight(root.left);
        }
        if (root.right != null) {
            rightHeight = getHeight(root.right);
        }

        return leftHeight - rightHeight;
    }

    /**
     * 获取某个节点的高度
     * 如果叶子节点高度是0，那没有叶子的一方高度应该是-1
     * 左右孩子高度最大值+1
     *
     * @param root
     * @return
     */
    public int getHeight(AVLNode<K> root) {
        if (root == null) {
            return -1;
        }
        if (root.left == null && root.right == null) {
            return 0;
        }

        return Math.max(getHeight(root.left), getHeight(root.right)) + 1;
    }


    /**
     * 中序遍历
     *
     * @param root
     */
    public void midOrder(AVLNode<K> root) {
        if (root == null) {
            return;
        }
        midOrder(root.left);
        System.out.println(root.key);
        midOrder(root.right);
    }

    /**
     * 先根遍历-递归实现
     *
     * @param t
     */
    public void preOrder(AVLNode<K> t) {
        //递归结束条件
        if (t == null) {
            return;
        } else {
            System.out.println(t.key);
        }
        preOrder(t.left);
        preOrder(t.right);

    }

    public static void main(String[] args) {
        BinaryAVLTree<Integer> avlTree = new BinaryAVLTree<>();
        List<Integer> ints = Lists.newArrayList(13, 12, 1, 5);

        for (Integer anInt : ints) {
            avlTree.insert(anInt);
        }
        System.out.println(ints);
        avlTree.remove(13);
        System.out.println("=======preOrder=========");
        avlTree.preOrder(avlTree.getRoot());
        System.out.println("=======midOrder=========");
        avlTree.midOrder(avlTree.getRoot());
        avlTree.insert(13);

        System.out.println("=======preOrder=========");
        avlTree.preOrder(avlTree.getRoot());
        System.out.println("=======midOrder=========");
        avlTree.midOrder(avlTree.getRoot());
        avlTree.remove(1);

        System.out.println("=======preOrder=========");
        avlTree.preOrder(avlTree.getRoot());
        System.out.println("=======midOrder=========");
        avlTree.midOrder(avlTree.getRoot());
    }
}