AVL树的Java实现

AVL树的Java实现

1. 前言

AVL树是带有平衡条件的二叉查找树，如果我们向一棵树插入排好序的一组数，树会变得很深而且左右两边不平衡。或者删除数据的时候也经常有这种情况，这时候插入和删除操作的复杂度会大幅上升。引入平衡条件后，就解决了这个问题。具体分析参见《数据结构与算法分析》第四版二叉搜索树和AVL树的章节。

2. 代码

package data.structure;

//节点定义
class AvlTreeNode {
    Integer value;
    AvlTreeNode left = null;
    AvlTreeNode right = null;
    Integer height = 0;

    @Override
    public String toString() {
        return "AvlTreeNode{" +
                "value=" + value +
                ", height=" + height +
                '}';
    }

    public AvlTreeNode(Integer value) {
        this.value = value;
    }
    public AvlTreeNode() {
    }
}

public class AvlTree {
    //获取树的高度，定义空节点的高度为-1，叶子节点的高度为0
    int getTreeHeight(AvlTreeNode root){
        return root == null ? -1 : root.height;
    }

    //插入数据和搜索树一样，多了个平衡操作
    AvlTreeNode insert(AvlTreeNode root, int value){
        if(root.value == null){
            root.value = value;
            return root;
        } else {
            if(root.value > value){
                if(root.left != null){
                    root.left = insert(root.left, value);
                } else {
                    root.left = new AvlTreeNode(value);
                }
            } else if(root.value < value) {
                if(root.right != null){
                    root.right = insert(root.right, value);
                } else {
                    root.right = new AvlTreeNode(value);

                }
            }
        }

        //插入完成后执行平衡操作
        root = balance(root);
        //返回平衡完成后的根节点
        return root;
    }

    /*
    * 不平衡的四种情况：
    * 1. 向左子树的左子树插入数据导致不平衡
    * 2. 向左子树的右子树插入数据导致不平衡
    * 3. 向右子树的左子树插入数据导致不平衡
    * 4. 向右子树的右子树插入数据导致不平衡
    * */
    AvlTreeNode balance(AvlTreeNode root){
        if(root == null){
            return null;
        }
        //树的两边不平衡，而且属于第一或第二种情况
        if(getTreeHeight(root.left) - getTreeHeight(root.right) > 1){
            //第一种情况，单旋转
            if(getTreeHeight(root.left.left) >= getTreeHeight(root.left.right)){
                root = rotateWithLeftChild(root);
            } else {
                //第二种情况，双旋转
                root = doubleRotateLeftChild(root);
            }

        } else if(getTreeHeight(root.right) - getTreeHeight(root.left) > 1){
            //第四种情况，单旋转
            if(getTreeHeight(root.right.right) > getTreeHeight(root.right.left)){
                root = rotateWithRightChild(root);
            } else {
                //第三种情况，双旋转
                root = doubleRotateRightChild(root);
            }
        }

        root.height = Integer.max(getTreeHeight(root.right),getTreeHeight(root.left)) + 1;
        return root;
    }

    AvlTreeNode rotateWithLeftChild(AvlTreeNode root){
        AvlTreeNode temp = root.left;
        //root.left的右结点值比root.left大，比root小，所以在root的左边
        root.left = temp.right;
        temp.right = root;
        //修改高度
        root.height = Integer.max(getTreeHeight(root.right),getTreeHeight(root.left)) + 1;
        temp.height = Integer.max(getTreeHeight(temp.left),root.height) + 1;

        //替换掉节点
        root = temp;
        //返回被替换后的根节点
        return root;
    }

    AvlTreeNode rotateWithRightChild(AvlTreeNode root){
        AvlTreeNode temp = root.right;
        root.right = temp.left;
        temp.left = root;
        root.height = Integer.max(getTreeHeight(root.right),getTreeHeight(root.left)) + 1;
        temp.height = Integer.max(getTreeHeight(temp.left),root.height) + 1;

        //替换掉节点
        root = temp;
        //返回被替换后的根节点
        return root;
    }
    //双旋转左子树
    AvlTreeNode doubleRotateLeftChild(AvlTreeNode root){
        //先将左子树的右节点左旋，例如{5,3,4},旋转完成后变成{5,4,3}
        root.left = rotateWithRightChild(root.left);
        //再单旋转一次，变成{4,3,5}
        root = rotateWithLeftChild(root);
        return root;
    }
    //双旋转右子树
    AvlTreeNode doubleRotateRightChild(AvlTreeNode root){
        root.right = rotateWithLeftChild(root.right);
        root = rotateWithRightChild(root);
        return root;
    }
    //先根遍历
    void print(AvlTreeNode root){
        if(root != null){
            System.out.print(root + " ");
            if(root.left != null){
                print(root.left);
            }
            if(root.right != null){
                print(root.right);
            }
        }
        return;
    }
    //寻找值最小的节点
    AvlTreeNode findMin(AvlTreeNode root){
        if(root == null){
            return null;
        }
        if(root.left!=null){
            return findMin(root.left);
        } else{
            return root;
        }
    }
    //删除操作和普通的搜索树一样，只不过多加了一个平衡
    AvlTreeNode delete(AvlTreeNode root, Integer value){
        if(root == null){
            return null;
        }
        if(root.value > value){
            root.left = delete(root.left, value);
        } else if(root.value < value){
            root.right = delete(root.right, value);
        } else if(root.right!=null && root.left!=null) {
            root.value = findMin(root.right).value;
            root.right = delete(root.right, root.value);
        } else{
            root =  root.left != null?root.left:root.right;
        }
        //删除完成后执行平衡操作
        root = balance(root);
        //返回平衡完成后的根节点
        return root;
    }

    void run(){
        int[] nums = {1,2,3,4,5,6,7,16,15,14,13,12,11,10,9,8};
        //int[] nums = {1,2,3};
        //int[] nums = {5,4,3,2,1};
        //int[] nums = {5,3,4};
        AvlTreeNode root = new AvlTreeNode();

        for (int num : nums) {

            root = insert(root, num);
            //插入一次输出一次
            print(root);
            System.out.println();
        }

        delete(root,4);
        print(root);
    }

    public static void main(String[] args) {
        AvlTree avlTree = new AvlTree();
        avlTree.run();

    }

}