AVL树 JAVA整理

总是忘记，还是要自己敲一遍才行。

---

• AVL树旋转：LL单旋、LR双旋、RR单旋、RL双旋
• AVL操作：insert、remove
• 二叉树排序：前序、中序、后序
• 插入，查找，删除的时间复杂度O(logN)。频繁旋转会使插入和删除牺牲掉O(logN)左右的时间
• 树形输出
• Test code
• 参考资料：见代码文档注释

---

二叉树排序中的前 中 后是指遍历时访问根节点的先后顺序。具体见代码。

---

package com.rookie.test;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;

/**
 * AvlTree
 * 
 *     AVL树JAVA：   https://www.cnblogs.com/skywang12345/p/3577479.html
 *     AVL树C语言：  https://blog.csdn.net/sjg_sjk/article/details/80332151
 *     AVL树插入图解：https://www.cnblogs.com/zhuwbox/p/3636783.html
 *     AVL删除图解：  https://blog.csdn.net/goodluckwhh/article/details/11786079
 *     二叉树排序：   https://blog.csdn.net/qq_33243189/article/details/80222629
 * 
 * @author: Mr.Rookie
 * @param 
 */
public class AvlTree> {

    /**
     * 树节点
     * @param 
     */
    public static class AvlTreeNode> {

        private D   data;

        private int height;

        private AvlTreeNode left;

        private AvlTreeNode right;


        public AvlTreeNode(D data, AvlTreeNode left, AvlTreeNode right) {
            this.data   = data;
            this.height = 1;
            this.left   = left;
            this.right  = right;
        }
    }


    /**
     * 当前树根节点
     */
    private AvlTreeNode rootNode;



    private int getHeight(AvlTreeNode tree) {
        //维基百科上：树的高度为最大层次。so,空的二叉树的高度是0
        if ( tree == null ) {
            return 0;
        }
        return tree.height;
    }


    /**
     * 当前树高
     * @return
     */
    public int getHeight() {
        return getHeight(rootNode);
    }


    /**
     * LL, 左左单旋
     * @param node 失衡节点，让其左孩子节点变成失衡节点的根节点
     * @return
     */
    private AvlTreeNode llRotation(AvlTreeNode node) {
        System.out.println("LL : " + node.data);

        // 原失衡节点的左孩子节点作为新的根节点
        AvlTreeNode newRoot = node.left;

        // 原失衡节点的左孩子节点替换为新根的右孩子节点
        node.left = newRoot.right;// 也既 node.left = node.left.right;
        // 新根节点的右孩子节点为原来的失衡节点
        newRoot.right = node;


        // 重新计算高度
        node.height    = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
        newRoot.height = Math.max(getHeight(newRoot.left), getHeight(newRoot.right)) + 1; // 现在 newRoot.right = node

        return newRoot;
    }



    /**
     * RR, 右右单旋
     * @param node 失衡节点，让其右孩子节点变成失衡节点的根节点
     * @return
     */
    private AvlTreeNode rrRotation(AvlTreeNode node) {
        System.out.println("RR : " + node.data);

        // 原失衡节点的右孩子节点作为新的根节点
        AvlTreeNode newRoot = node.right;

        // 原失衡节点的右孩子节点替换为新根的左孩子节点
        node.right = newRoot.left;// 也既 node.right = node.right.left;
        // 新根节点的左孩子节点为原来的失衡节点
        newRoot.left = node;


        // 重新计算高度
        node.height    = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
        newRoot.height = Math.max(getHeight(newRoot.left), getHeight(newRoot.right)) + 1; // 现在 newRoot.left = node

        return newRoot;
    }


    /**
     * LR  左右情况， 两次旋转
     * @param node
     * @return
     */
    private AvlTreeNode lrRotation(AvlTreeNode node) {
        System.out.println("LR : " + node.data);

        // 1 对失衡节点的左孩子节点进行rr单旋
        node.left = rrRotation(node.left);
        // 2 再对失衡节点进行ll单旋
        return llRotation(node);
    }


    /**
     * RL  右左情况， 两次旋转
     * @param node
     * @return
     */
    private AvlTreeNode rlRotation(AvlTreeNode node) {
        System.out.println("RL : " + node.data);

        // 1 对失衡节点的右孩子节点进行ll单旋
        node.right = llRotation(node.right);
        // 2 再对失衡节点进行ll单旋
        return rrRotation(node);
    }


    /**
     * 向树中插入数据
     * @param data
     */
    public void insert(T data) {
        if ( data == null ) {
            return;
        }
        rootNode = insert(rootNode, data);
    }


    /**
     * 插入操作
     * @param treeRoot 需要插入的树根节点
     * @param data 需要插入的节点数据
     * @return
     */
    private AvlTreeNode insert(AvlTreeNode treeRoot, T data) {

        if ( treeRoot == null ) {
            // 新建节点， 高度为1
            treeRoot = new AvlTreeNode<>(data, null, null);
            return treeRoot;
        }

        int comp = data.compareTo(treeRoot.data);
        if ( comp == 0 ) {
            // 不允许添加相同的节点值
            return treeRoot;

        } else if ( comp < 0 ) {

            // 放到treeRoot左子树
            treeRoot.left = insert(treeRoot.left, data);

            // 判断是否失衡 左子树高度 - 右子树高度   （data是被插入到左子树，所以只会引起左子树的高度，不变或增高）
            if ( getHeight(treeRoot.left) - getHeight(treeRoot.right) > 1 ) {
                // 左右节点高度差大于1，已经失衡，判断属于 LL， 还是属于LR
                if ( data.compareTo(treeRoot.left.data) < 0 ) {
                    // 插入的节点和当前失衡节点的左孩子节点比较，若节点插入到失衡节点的左孩子节点的左子树中， 属于LL情况
                    treeRoot = llRotation(treeRoot);
                } else {
                    // 插入的节点在失衡节点的左孩子节点的右子树中，属于LR情况
                    treeRoot = lrRotation(treeRoot);
                }
            }
        } else { // comp > 0

            // 放到treeRoot的右子树
            treeRoot.right = insert(treeRoot.right, data);

            // 判断是否失衡 右子树高度 - 左子树高度   （data是被插入到右子树，所以只会引起右边子树的高度，不变或增高）
            if ( getHeight(treeRoot.right) - getHeight(treeRoot.left) > 1 ) {
                // 判断属于 RR，还是属于RL情况
                if ( data.compareTo(treeRoot.right.data) > 0 ) {
                    // 当前节点被插入到失衡节点的右孩子节点的右子树中，则为RR
                    treeRoot = rrRotation(treeRoot);
                } else {
                    // 被插入到失衡节点的有孩子节点的左子树中，则为RL
                    treeRoot = rlRotation(treeRoot);
                }
            }
        }
        treeRoot.height = Math.max(getHeight(treeRoot.left), getHeight(treeRoot.right)) + 1;
        return treeRoot;
    }


    /**
     * 从树中移除数据
     * @param data
     */
    public void remove(T data) {
        if ( data == null ) {
            return;
        }
        this.rootNode = remove(this.rootNode, data);
    }


    /**
     * 从树中移除包含值的节点
     * @param tree
     * @param data
     * @return
     */
    private AvlTreeNode remove(AvlTreeNode tree, T data) {

        if ( tree == null || data == null ) {
            return tree;
        }

        int comp = data.compareTo(tree.data);
        if ( comp < 0 ) {
            // 要删除的节点在tree的左子树中，继续递归删除
            tree.left = remove(tree.left, data);
            // 检查平衡性
            if (getHeight(tree.right) - getHeight(tree.left) > 1) {
                // 当前节点的右子树高于其左子树
                // 检查当前节点的右子树的左右子树高度差，判断旋转情况
                if ( getHeight(tree.right.left) > getHeight(tree.right.right) ) {
                    // 右左情况 RL双旋
                    tree = rlRotation(tree);
                } else {
                    // 右右情况 RR单旋
                    tree = rrRotation(tree);
                }
            }
        } else if (comp > 0 ) {
            // 要删除的节点在tree的有子树中，继续递归删除
            tree.right = remove(tree.right, data);
            // 删除之后需要检查当前树的平衡
            if ( getHeight(tree.left) - getHeight(tree.right) > 1 ) {
                // 左子树比右子树高
                // 检查左子树的左子树高度 和 左子树的右子树高度，判断旋转情况
                if ( getHeight(tree.left.right) > getHeight(tree.left.left) ) {
                    // 左右情况， LR双旋
                    tree = lrRotation(tree);
                } else {
                    // 左左情况。 LL单旋
                    tree = llRotation(tree);
                }
            }
        } else {
            // 当前节点是需要删除的节点
            if ( tree.left == null || tree.right == null ) {
                // 要删除的节点没有子树，或只有一个子树
                // 将要删除的节点替换为其子树（若有子树替换，没有则为null）
                tree = (tree.left != null) ? tree.left : tree.right;
            } else {
                // 要删除的节点 左右子树都有， 都不为空
                // 比较要删除的节点的左右子树高度
                if ( getHeight(tree.left) > getHeight(tree.right) ) {
                    // 删除节点的左子树高于右子树，
                    // 1 找到左子树中最大的节点
                    AvlTreeNode maxChild = maxTreeNode(tree.left);
                    // 2 将删除节点的值替换为其左子树中的最大值 （只替换值，不用移动节点，避免移动指针和重新计算高度）
                    tree.data = maxChild.data;
                    // 3 替换后从左子树中删除最大值；删除之后当前子树还是保持平衡
                    tree.left = remove(tree.left, maxChild.data);
                } else {
                    // 删除的节点 左右子树等高，或右子树高
                    // 1 从删除节点的右子树中找到最小的节点
                    AvlTreeNode minChild = minTreeNode(tree.right);
                    // 2 将最小值节点值替换到当前删除节点中。（只需要替换值，原因同上）
                    tree.data = minChild.data;
                    // 3 替换后从右子树中删除最小值；删除之后当前子树还是保持平衡
                    tree.right = remove(tree.right, minChild.data);
                }
            }
        }
        return tree;
    }


    /**
     * 返回子树中最大的节点 （非递归实现）
     * @param tree
     * @return
     */
    private AvlTreeNode maxTreeNode(AvlTreeNode tree) {

        if ( tree == null ) {
            return null;
        }

        AvlTreeNode node = tree;
        // 一直搜索当前节点的右节点
        while ( node.right != null ) {
            node = node.right;
        }

        return node;
    }


    /**
     * 返回子树中最小的节点（递归实现）
     * @param tree notNull
     * @return
     */
    private AvlTreeNode minTreeNode(AvlTreeNode tree) {
        if ( tree == null || tree.left == null ) {
            return tree;
        }
        return minTreeNode(tree.left);
    }


    /**
     * 返回当前树最大值
     * @return
     */
    public T maxData() {
        if ( rootNode == null ) {
            return null;
        }
        return maxTreeNode(this.rootNode).data;
    }


    /**
     * 返回当前树的最小值
     * @return
     */
    public T minData() {
        if ( rootNode == null ) {
            return null;
        }
        return minTreeNode(this.rootNode).data;
    }



    /**
     * 前序遍历
     * 
     *     先中，再左，最后右
     * 
     * @param container
     * @param tree
     */
    private void preorderTraversal(List container, AvlTreeNode tree) {
        if ( tree == null ) {
            return;
        }
        container.add(tree.data);
        preorderTraversal(container, tree.left);
        preorderTraversal(container, tree.right);
    }

    /**
     * 中序遍历
     * 
     *     先左，再中，最后右
     * 
     * @param container
     * @param tree
     */
    private void inorderTraversal(List container, AvlTreeNode tree) {
        if ( tree == null ) {
            return;
        }
        inorderTraversal(container, tree.left);
        container.add(tree.data);
        inorderTraversal(container, tree.right);
    }

    /**
     * 后序遍历
     * 
     *     先左，再右，最后中
     * 
     * @param container
     * @param tree
     */
    private void postorderTraversal(List container, AvlTreeNode tree) {
        if ( tree == null ) {
            return;
        }
        postorderTraversal(container, tree.left);
        postorderTraversal(container, tree.right);
        container.add(tree.data);
    }

    /**
     * 前序遍历结果列表
     * @return
     */
    public List preorderTraversal() {
        if ( this.rootNode == null ) {
            return Collections.emptyList();
        }

        List result = new ArrayList<>();
        preorderTraversal(result, this.rootNode);
        return result;
    }

    /**
     * 中序遍历
     * @return
     */
    public List inorderTraversal() {
        if ( this.rootNode == null ) {
            return Collections.emptyList();
        }

        List result = new ArrayList<>();
        inorderTraversal(result, this.rootNode);
        return result;
    }


    /**
     * 后序遍历
     * @return
     */
    public List postorderTraversal() {
        if ( this.rootNode == null ) {
            return Collections.emptyList();
        }

        List result = new ArrayList<>();
        postorderTraversal(result, this.rootNode);
        return result;
    }


// ******************************************************************************

    public static void main(String[] args) {
//        testAddAvlTree();
        testRemove();
    }


    public static void testAddAvlTree() {
        AvlTree tree = new AvlTree<>();
        Random random = ThreadLocalRandom.current();

        int num;
        for (int i = 0; i<32; i++) {
            num = random.nextInt(32);
            System.out.println("----> add: " + num);

            tree.insert(num);

            System.out.println("----> result : ");
            tree.printTree();

            System.out.println("");
            System.out.println("==================================");
        }
    }


    public static void testRemove() {

        int[] nums = new int[]{17,12,6,3,0,4,8,10,15,13,16,28,23,18,25,24,26,30,29,31};

        AvlTree tree = new AvlTree<>();
        for (int num : nums) {
            tree.insert(num);
        }

        tree.printTree();

        System.out.println("==================================");
        int remove = 6;
        System.out.println("remove: " + remove);
        tree.remove(remove);
        tree.printTree();

        System.out.println("==================================");
        remove = 10;
        System.out.println("remove: " + remove);
        tree.remove(remove);
        tree.printTree();

        System.out.println("==================================");
        remove = 25;
        System.out.println("remove: " + remove);
        tree.remove(remove);
        tree.printTree();
    }


// ******************************************************************************

    /**
     * 以树形方式输出树 (此方法主要针对数值类型的树)
     */
    public void printTree() {

        if ( this.rootNode == null ) {
            System.out.println("None...");
            return;
        }

        // 只输出数值类型 - 便于测试
        if ( !Number.class.isInstance(this.rootNode.data) ) {
            System.out.println("Only for Number types ...");
            return;
        }

        StringBuilder buff = new StringBuilder();
        int charLen = maxData().toString().length();// 最大数据长度
        int initGap = charLen; // 填充的空格单元长度

        List rowContainer = new ArrayList<>();
        int nullNum = 0; // 每一行开始需要天正的空格单元数量
        int gapNum = 0;  // 每一行数据间隔的单元数量
        StringBuilder row= new StringBuilder();

        for ( int h = 1; h <= getHeight(); h++) {
            rowContainer.clear();
            getDatasByLayer(rowContainer, rootNode, h, getHeight(rootNode));

            addNullSpace(row, nullNum * initGap);
            gapNum = 2 * nullNum + 1;
            nullNum = gapNum;
            for (T t : rowContainer) {
                addData(row, t, charLen);
                addNullSpace(row, gapNum * initGap);
            }
            buff.insert(0, row);
            row.setLength(0);

            if ( h != getHeight() ) {
                buff.insert(0, "\r\n");
                buff.insert(0, "\r\n");
            }
        }

        buff.append("\r\n");
        buff.append("\r\n");

        // 前序
        rowContainer.clear();
        preorderTraversal(rowContainer, this.rootNode);
        buff.append("前序：").append(rowContainer.stream().map(String::valueOf).collect(Collectors.joining(","))).append("\r\n");

        // 中序
        rowContainer.clear();
        inorderTraversal(rowContainer, this.rootNode);
        buff.append("中序：").append(rowContainer.stream().map(String::valueOf).collect(Collectors.joining(","))).append("\r\n");

        // 后序
        rowContainer.clear();
        postorderTraversal(rowContainer, this.rootNode);
        buff.append("后序：").append(rowContainer.stream().map(String::valueOf).collect(Collectors.joining(","))).append("\r\n");
        rowContainer.clear();

        System.out.println(buff);
        buff.setLength(0);
    }

    /**
     * 填充空格字符
     * @param buff
     * @param len
     */
    private void addNullSpace(StringBuilder buff, int len) {
        if (len <=0) {
            return;
        }
        for (int i=1; i<=len; i++) {
            buff.append(" ");
        }
    }

    /**
     * 填充数据
     * @param buff
     * @param t
     * @param len
     */
    private void addData(StringBuilder buff, T t, int len) {
        String s = t == null ? "X" : t.toString();
        if ( s.length() < len ) {
            addNullSpace(buff, len - s.length());
        }
        buff.append(s);
    }

    /**
     * 返回指定层数的所有数据
     * @param layer
     * @return
     */
    private void getDatasByLayer(List container, AvlTreeNode tree, int layer, int currentLayer) {

        if ( currentLayer == layer) {
            container.add(tree == null ? null : tree.data);
            return;
        }

        currentLayer -= 1;
        if ( currentLayer < layer ) {
            return;
        }

        if ( tree == null ) {
            int size = (int)Math.pow(2, currentLayer - layer + 1);
            for (int i=1; i<= size; i++) {
                container.add(null);
            }
            return;
        }

        getDatasByLayer(container, tree.left,  layer, currentLayer);
        getDatasByLayer(container, tree.right, layer, currentLayer);
    }


}