1
2
3
4
5
6
7
8
9

public class TreeNode {
    int val;
    TreeNode left, right;
    int height;

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

1
2
3
4
5
6
7
8

public int height(TreeNode root) {
    return root == null ? -1 : root.height;
}

public void updateHeight(TreeNode root) {
    if (root == null) return;
    root.height = Math.max(height(root.left), height(root.right)) + 1;
}

1
2
3
4

public int balanceFactor(TreeNode root) {
    if (root == null) return 0;
    return height(root.left) - height(root.right);
}

1
2
3
4
5
6
7
8

public TreeNode rightRotate(TreeNode root) {
    TreeNode child = root.left, grandChild = child.right;
    child.right = root;
    root.left = grandChild;
    updateHeight(root);
    updateHeight(child);
    return child;
}

1
2
3
4
5
6
7
8

public TreeNode leftRotate(TreeNode root) {
    TreeNode child = root.right, grandChild = child.left;
    child.left = root;
    root.right = grandChild;
    updateHeight(root);
    updateHeight(child);
    return child;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

public TreeNode rotate(TreeNode root) {
    int balanceFactor = balanceFactor(root);
    // 左偏树
    if (balanceFactor > 1) {
        if (balanceFactor(root.left) >= 0) return rightRotate(root); // LL
        else {  // LR
            root.left = leftRotate(root.left);
            return rightRotate(root);
        }
    } else if (balanceFactor < -1) { // 右偏树
        if (balanceFactor(root.right) <= 0) return leftRotate(root); // RR
        else {  // RL
            root.right = rightRotate(root.right);
            return leftRotate(root);
        }
    }
    // 其他情况是平衡树，无需旋转
    return root;
}

1
2
3
4
5
6
7
8
9

public TreeNode insert(TreeNode root, int val) {
    if (root == null) return new TreeNode(val);
    if (val < root.val) root.left = insert(root.left, val);
    else if (val > root.val) root.right = insert(root.right, val);
    else return root; // 重复节点，不插入
    updateHeight(root);  // 插入后更新高度
    root = rotate(root); // 旋转操作，使子树平衡
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

public TreeNode insert(TreeNode root, int val) {
    if (root == null) return new TreeNode(val);
    TreeNode cur = root, pre = null;
    Deque<TreeNode> stack = new ArrayDeque<>();
    // 找到插入位置
    while (cur != null) {
        stack.push(cur);
        pre = cur;
        if (val == cur.val) return root; // 重复值, 不插入
        if (val < cur.val) cur = cur.left;
        else cur = cur.right;
    }
    // 插入节点
    if (val < pre.val) pre.left = new TreeNode(val);
    else pre.right = new TreeNode(val);
    // 更新高度, 旋转维护 AVL 平衡, 自底向上
    while (!stack.isEmpty()) {
        TreeNode node = stack.pop();
        updateHeight(node);
        node = rotate(node);
        // 旋转后, 子树根节点可能变化, 更新父节点的子节点指向
        if (!stack.isEmpty()) {
            TreeNode parent = stack.peek(); // 不能 pop, 下一次 while-loop 需要更新高度和旋转
            if (node.val < parent.val) parent.left = node;
            else parent.right = node;
        } else {
            root = node;  // node 即是 AVL 的根节点
        }
    }
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

public TreeNode remove(TreeNode root, int val) {
    if (root == null) return null;
    // 找节点删除
    if (val < root.val) root.left = remove(root.left, val);
    else if (val > root.val) root.right = remove(root.right, val);
    else {
        // 子节点有 0 或 1 个
        if (root.left == null || root.right == null) {
            TreeNode child = root.left == null ? root.right : root.left;
            if (child == null) root = null; // 无子节点，整树删除
            else root = child;  // 有一个子节点，子节点为新根
        } else { // 子节点有 2 个
            TreeNode successor = root.right;
            while (successor.left != null) successor = successor.left;
            root.right = remove(root.right, successor.val);
            root.val = successor.val;  // 值替换
        }
    }
    updateHeight(root); // 删除后更新高度
    root = rotate(root); // 旋转操作，使子树平衡
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

public TreeNode remove(TreeNode root, int val) {
    if (root == null) return null;
    // 找节点删除
    if (val < root.val) root.left = remove(root.left, val);
    else if (val > root.val) root.right = remove(root.right, val);
    else {
        // 子节点有 0 或 1 个
        if (root.left == null || root.right == null) {
            TreeNode child = root.left == null ? root.right : root.left;
            if (child == null) root = null; // 无子节点，整树删除
            else root = child;  // 有一个子节点，子节点为新根
        } else { // 子节点有 2 个
            TreeNode successor = root.right;
            while (successor.left != null) successor = successor.left; // 找到右子树的最左节点
            // 后继节点替换 root; 先将其从右子树删除, 再直接替代 root, 最后需要把 root 更新为后继节点
            root.right = remove(root.right, successor.val);
            successor.left = root.left;
            successor.right = root.right;
            root = successor;
        }
    }
    updateHeight(root); // 删除后更新高度
    root = rotate(root); // 旋转操作，使子树平衡
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

public TreeNode remove(TreeNode root, int val) {
    if (root == null) return null;
    // 找节点删除
    TreeNode cur = root, pre = null;
    Deque<TreeNode> stack = new ArrayDeque<>();
    while (cur != null) {
        stack.push(cur); // 记录路径
        if (val == cur.val) break;
        pre = cur;
        if (val < cur.val) cur = cur.left;
        else cur = cur.right;
    }
    if (cur == null) return root;  // 没有找到, 忽略删除
    // cur 就是待删除节点
    if (cur.left == null || cur.right == null) { // 待删除节点有 0 或 1 个子节点
        TreeNode child = cur.left == null ? cur.right : cur.left;
        // 删除根节点, 不影响子树高度和平衡, 可直接返回
        if (pre == null) return child;
        if (cur == pre.left) pre.left = child;
        else pre.right = child;
    } else {  // 有两个子节点
        TreeNode successor = cur.right, successorParent = cur;
        while (successor.left != null) {
            stack.push(successor);  // 记录路径
            successorParent = successor;
            successor = successor.left;
        }
        // cur 右子树删除 successor 节点
        if (successorParent == cur) {  // successor 是 cur 的直接右子节点
            successorParent.right = successor.right;
        } else {
            successorParent.left = successor.right;
        }
        // 删除 cur, successor 替换 cur
        successor.left = cur.left;
        successor.right = cur.right;
        // 维护新的 root(cur 的父节点)
        if (pre == null) root = successor; // 删除根节点
        else if (cur == pre.left) pre.left = successor;
        else pre.right = successor;
    }
    // 更新高度, 旋转维护 AVL 平衡, 自底向上
    while (!stack.isEmpty()) {
        TreeNode node = stack.pop();
        updateHeight(node);
        node = rotate(node);
        // 旋转后, 子树根节点可能变化, 更新父节点的子节点指向
        if (!stack.isEmpty()) {
            TreeNode parent = stack.peek(); // 不能 pop, 下一次 while-loop 需要更新高度和旋转
            if (node.val < parent.val) parent.left = node;
            else parent.right = node;
        } else {
            root = node;  // node 即是 AVL 的根节点
        }
    }
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

public TreeNode remove(TreeNode root, int val) {
    if (root == null) return null;
    // 找节点删除
    TreeNode cur = root, pre = null;
    Deque<TreeNode> stackCur = new ArrayDeque<>();  // 记录 Cur 之前的路径
    Deque<TreeNode> stackSucc = new ArrayDeque<>(); // 记录 Cur 到 Successor 的路径
    while (cur != null) {
        stackCur.push(cur); // 记录路径
        if (val == cur.val) break;
        pre = cur;
        if (val < cur.val) cur = cur.left;
        else cur = cur.right;
    }
    if (cur == null) return root;  // 没有找到, 忽略删除
    stackCur.pop(); // cur 会被删除, 所以不需要再更新
    // cur 就是待删除节点
    if (cur.left == null || cur.right == null) { // 待删除节点有 0 或 1 个子节点
        TreeNode child = cur.left == null ? cur.right : cur.left;
        // 删除根节点, 不影响子树高度和平衡, 可直接返回
        if (pre == null) return child;
        if (cur == pre.left) pre.left = child;
        else pre.right = child;
    } else {  // 有两个子节点
        TreeNode successor = cur.right, successorParent = cur;
        while (successor.left != null) {
            stackSucc.push(successor);  // 记录路径
            successorParent = successor;
            successor = successor.left;
        }
        // cur 右子树删除 successor 节点
        if (successorParent == cur) {  // successor 是 cur 的直接右子节点
            successorParent.right = successor.right;
        } else {
            successorParent.left = successor.right;
        }
        // 删除 cur, successor 替换 cur
        successor.left = cur.left;
        successor.right = cur.right;
        stackCur.push(successor);  // successor 替换 cur, 加入到 stackCur 中, 其父节点时 pre
        // 维护新的 root(cur 的父节点)
        if (pre == null) root = successor; // 删除根节点
        else if (cur == pre.left) pre.left = successor;
        else pre.right = successor;
    }
    // 更新高度, 旋转维护 AVL 平衡, 自底向上, 所以从 stackSucc 开始
    while (!stackSucc.isEmpty()) {
        TreeNode node = stackSucc.pop();
        updateHeight(node);
        node = rotate(node);
        // 旋转后, 子树根节点可能变化, 更新父节点的子节点指向
        if (!stackSucc.isEmpty()) {
            TreeNode parent = stackSucc.peek(); // 不能 pop, 下一次 while-loop 需要更新高度和旋转
            if (node.val < parent.val) parent.left = node;
            else parent.right = node;
        }
    }
    // 然后再更新 stackCur, 区别于 stackSucc, 多了 root 节点的更新, 可能删除的节点是 root
    while (!stackCur.isEmpty()) {
        TreeNode node = stackCur.pop();
        updateHeight(node);
        node = rotate(node);
        // 旋转后, 子树根节点可能变化, 更新父节点的子节点指向
        if (!stackCur.isEmpty()) {
            TreeNode parent = stackCur.peek(); // 不能 pop, 下一次 while-loop 需要更新高度和旋转
            if (node.val < parent.val) parent.left = node;
            else parent.right = node;
        } else {
            root = node;  // node 即是 AVL 的根节点
        }
    }
    return root;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

public TreeNode remove(TreeNode root, int val) {
    if (root == null) return null;
    // 找节点删除
    TreeNode cur = root, pre = null;
    Deque<TreeNode> stack = new ArrayDeque<>();
    while (cur != null) {
        stack.push(cur); // 记录路径
        if (val == cur.val) break;
        pre = cur;
        if (val < cur.val) cur = cur.left;
        else cur = cur.right;
    }
    if (cur == null) return root;  // 没有找到, 忽略删除
    // cur 就是待删除节点
    if (cur.left == null || cur.right == null) { // 待删除节点有 0 或 1 个子节点
        TreeNode child = cur.left == null ? cur.right : cur.left;
        // 删除根节点, 不影响子树高度和平衡, 可直接返回
        if (pre == null) return child;
        if (cur == pre.left) pre.left = child;
        else pre.right = child;
        if (child != null)
            replaceStack(stack, cur, child); // 将栈中 cur 替换为 child
        else
            stack.pop();  // 栈顶 cur 已被删除, 且没有新的子节点入栈
    } else {  // 有两个子节点
        TreeNode successor = cur.right, successorParent = cur;
        while (successor.left != null) {
            stack.push(successor);  // 记录路径
            successorParent = successor;
            successor = successor.left;
        }
        // cur 右子树删除 successor 节点
        if (successorParent == cur) {  // successor 是 cur 的直接右子节点
            successorParent.right = successor.right;
        } else {
            successorParent.left = successor.right;
        }
        // 删除 cur, successor 替换 cur
        successor.left = cur.left;
        successor.right = cur.right;
        replaceStack(stack, cur, successor); // 将栈中 cur 替换为 successor
        // 维护新的 root(cur 的父节点)
        if (pre == null) root = successor; // 删除根节点
        else if (cur == pre.left) pre.left = successor;
        else pre.right = successor;
    }
    // 更新高度, 旋转维护 AVL 平衡, 自底向上
    while (!stack.isEmpty()) {
        TreeNode node = stack.pop();
        updateHeight(node);
        node = rotate(node);
        // 旋转后, 子树根节点可能变化, 更新父节点的子节点指向
        if (!stack.isEmpty()) {
            TreeNode parent = stack.peek(); // 不能 pop, 下一次 while-loop 需要更新高度和旋转
            if (node.val < parent.val) parent.left = node;
            else parent.right = node;
        } else {
            root = node;  // node 即是 AVL 的根节点
        }
    }
    return root;
}

// 将 stack 中 oldNode 替换为 newNode
public void replaceStack(Deque<TreeNode> stack, TreeNode oldNode, TreeNode newNode) {
    Deque<TreeNode> tempStack = new LinkedList<>();
    while (!stack.isEmpty()) {
        TreeNode node = stack.pop();
        if (node == oldNode) {
            stack.push(newNode);
            break;
        } else {
            tempStack.push(node);
        }
    }
    // 将临时栈中的元素重新压入原栈
    while (!tempStack.isEmpty()) {
        stack.push(tempStack.pop());
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

/**
 * BST 查找
 *
 * @param root 根节点
 * @param num  待查找的值
 * @return 查找到的节点
 */
public TreeNode search(TreeNode root, int num) {
    TreeNode cur = root;
    while (cur != null) {
        if (num == cur.val) break;
        else if (num < cur.val) cur = cur.left;
        else cur = cur.right;
    }
    return cur;
}