【Java】平衡二叉树的实现代码

Node.java

public class Node<T extends Comparable<T>> {
    public T value;
    public Node<T> lChild;
    public Node<T> rChild;
    private int height;

    public Node(T value) {
        this.value = value;
        this.height = 1;
    }

    /**
     * 获得某个节点的高度
     */
    private int getHeight(Node<T> node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 计算左右子树的高度差
     */
    public int calcHeightDiff() {
        int lHeight = getHeight(this.lChild);
        int rHeight = getHeight(this.rChild);
        return (lHeight - rHeight);
    }

    /**
     * 判断当前节点是否左右平衡
     */
    public boolean isBalance() {
        int diff = this.calcHeightDiff();
        return (Math.abs(diff) <= 1);
    }

    /**
     * 记录当前节点的高度
     */
    public void refresh() {
        int lHeight = getHeight(this.lChild);
        int rHeight = getHeight(this.rChild);
        this.height = Math.max(lHeight, rHeight) + 1;
    }

}

AVLTree.java

public class AVLTree<T extends Comparable<T>> {
    private Node<T> root;

    /**
     * 添加指定值到二叉平衡树中
     */
    public void add(T value) {
        Node<T> newNode = new Node<>(value);
        this.root = setNode(this.root, newNode);
    }

    /**
     * 根据值将节点放置特定位置
     */
    private Node<T> setNode(Node<T> curNode, Node<T> newNode) {
        if (curNode == null) {
            return newNode;
        }
        if (curNode.value.compareTo(newNode.value) > 0) { // put left
            curNode.lChild = setNode(curNode.lChild, newNode);
        } else if (curNode.value.compareTo(newNode.value) <= 0) { // put right
            curNode.rChild = setNode(curNode.rChild, newNode);
        }
        return adjustBalance(curNode);
    }

    /**
     * 调整二叉树节点的平衡状态
     */
    private Node<T> adjustBalance(Node<T> node) {
        if (node == null) {
            return null;
        }
        if (node.isBalance()) {
            node.refresh();
            return node;
        }
        if (node.calcHeightDiff() > 0) { // L?
            if (node.lChild.calcHeightDiff() >= 0) { // LL
                return rotateLL(node);
            } else { // LR
                return rotateLR(node);
            }
        } else if (node.calcHeightDiff() < 0) { // R?
            if (node.rChild.calcHeightDiff() <= 0) { // RR
                return rotateRR(node);
            } else { // RL
                return rotateRL(node);
            }
        }
        return null;
    }

    /**
     * 以节点为轴向右旋转
     */
    private Node<T> rotateLL(Node<T> nodeA) {
        Node<T> nodeB = nodeA.lChild;
        nodeA.lChild = nodeB.rChild;
        nodeB.rChild = nodeA;
        nodeA.refresh();
        nodeB.refresh();
        return nodeB;
    }

    /**
     * 以节点为轴向左旋转
     */
    private Node<T> rotateRR(Node<T> nodeA) {
        Node<T> nodeB = nodeA.rChild;
        nodeA.rChild = nodeB.lChild;
        nodeB.lChild = nodeA;
        nodeA.refresh();
        nodeB.refresh();
        return nodeB;
    }

    /**
     * 先子节点左旋后右旋
     */
    private Node<T> rotateLR(Node<T> nodeA) {
        nodeA.lChild = rotateRR(nodeA.lChild);
        return rotateLL(nodeA);
    }

    /**
     * 先子节点右旋后左旋
     */
    private Node<T> rotateRL(Node<T> nodeA) {
        nodeA.rChild = rotateLL(nodeA.rChild);
        return rotateRR(nodeA);
    }

    /**
     * 从二叉平衡树中删除指定值
     */
    public void remove(T value) {
        this.root = resetNode(this.root, value);
    }

    /**
     * 根据值重置节点所在的子树
     */
    private Node<T> resetNode(Node<T> node, T value) {
        if (node == null) {
            return null;
        }
        if (node.value.compareTo(value) == 0) {
            if (node.lChild == null && node.rChild == null) {
                return null;
            }
            if (node.lChild != null) {
                node.lChild = replaceRR(node.lChild, node);
            } else {
                node.rChild = replaceLL(node.rChild, node);
            }
        } else if (node.value.compareTo(value) > 0) {
            node.lChild = resetNode(node.lChild, value);
        } else if (node.value.compareTo(value) < 0) {
            node.rChild = resetNode(node.rChild, value);
        }
        return adjustBalance(node);
    }

    /**
     * 查找子树中的最大值做替换
     */
    private Node<T> replaceRR(Node<T> curNode, Node<T> repNode) {
        if (curNode.rChild == null) {
            repNode.value = curNode.value;
            if (curNode.lChild == null) {
                return null;
            } else {
                return curNode.lChild;
            }
        } else {
            curNode.rChild = replaceRR(curNode.rChild, repNode);
            return adjustBalance(curNode);
        }
    }

    /**
     * 查找子树中的最小值做替换
     */
    private Node<T> replaceLL(Node<T> curNode, Node<T> repNode) {
        if (curNode.lChild == null) {
            repNode.value = curNode.value;
            if (curNode.rChild == null) {
                return null;
            } else {
                return curNode.rChild;
            }
        } else {
            curNode.lChild = replaceLL(curNode.lChild, repNode);
            return adjustBalance(curNode);
        }
    }


    /**
     * 递归返回每个节点的字符串形式
     */
    private String toString(Node<T> node) {
        if (node == null) {
            return "";
        }
        if (node.lChild == null && node.rChild == null) {
            return node.value.toString();
        }
        String lStr = toString(node.lChild);
        String rStr = toString(node.rChild);
        return node.value + "(" + lStr + "," + rStr + ")";
    }

    @Override
    public String toString() {
        return toString(this.root);
    }

}

TestDemo.java

public class TestDemo {
    public static void main(String[] args) {
        // 随机生成将要被添加的数列
        List<Integer> list = new ArrayList<>();
        for (int i = 0; i < 50; i++) {
            list.add((int) (Math.random() * 10));
        }
        System.out.println(list);

        AVLTree<Integer> tree = new AVLTree<>();
        for (Integer i : list) {
            tree.add(i);
        }
        System.out.println(tree);
        // 随机生成将要被删除的数列
        list.clear();
        for (int i = 0; i < 50; i++) {
            list.add((int) (Math.random() * 10));
        }
        System.out.println(list);

        for (Integer i : list) {
            tree.remove(i);
        }
        System.out.println(tree);
    }

}

AVL Tree Visualzation 平衡二叉树树可视化
https://www.cs.usfca.edu/~galles/visualization/AVLtree.html