JavaScript实现AVL树

发表于 更新于 分类于 Waline：

最近尝试实现了下AVL树，发现这玩意写起来还是很多细节的，搞了半天，所以就在这里记录下。

AVL树就是平衡二叉检索树，一方面它是BST，即二叉检索树，一方面它是平衡的，也就是任何一个节点为根结点的子树的左子树和右子树的高度差不到1。

基本原理

关于BST就不多说了，比较简单，就是每个节点的左孩子都小于根节点，右孩子都大于根节点，插入的时候注意就好。

但是BST有个问题，比如依次插入1，2，3，4，5，6，那么最终会得到一个只有右孩子的树，其实这个BST已经退化成了普通链表，所以我们需要一定的方式保持这个树的平衡，而保持平衡的方式如下：

具体实现

四种不平衡的情况

左左情况

这个情况下要右旋，然后返回旋转后新的根节点，代码如下：

1
2
3
4
5
6
#_rotateRight(node) {
    const avlNode = node.left;
    node.left = avlNode.right;
    avlNode.right = node;
    return avlNode;
}

右右情况

这个情况下要左旋，然后返回旋转后新的根节点，代码如下：

1
2
3
4
5
6
#_rotateLeft(node) {
    const avlNode = node.right;
    node.right = avlNode.left;
    avlNode.left = node;
    return avlNode;
}

左右情况

这个情况下先要对根节点的左子节点做左旋，变为左左情况，然后再对根节点做右旋，然后返回旋转后新的根节点，代码如下：

1
2
3
4
#_rotateLeftRight(node) {
    node.left = this.#_rotateLeft(node.left);
    return this.#_rotateRight(node)
}

右左情况

这个情况下先要对根节点的右子节点做右旋，变为右右情况，然后再对根节点做左旋，然后返回旋转后新的根节点，代码如下：

1
2
3
4
#_rotateRightLeft(node) {
    node.right = this.#_rotateRight(node.right);
    return this.#_rotateLeft(node);
}

修复不平衡的情况

获取节点的高度

我们在做平衡的时候需要判断是否需要平衡以及是哪种不平衡的情况好选择不同的旋转方式

1
2
3
4
5
6
7
8
9
10
11
12
13
#_getAvlTreeHeight(node) {
    // 空节点的高度为0
    if (node === null) {
        return 0;
    }

    // 这里使用后序遍历，因为每次递归的结果需要子节点递归的结果
    const leftHeight = this.#_getAvlTreeHeight(node.left);
    const rightHeight = this.#_getAvlTreeHeight(node.right);

    // 节点高度为左右子树高度中的较大值加上1
    return Math.max(leftHeight, rightHeight) + 1;
}

平衡

平衡以node为根节点的子树，并返回新的根节点。

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
#_balance(node) {
    // 空节点不需要平衡，直接返回空节点
    if (node === null) {
        return node;
    }

    // 获取左右子树的高度
    const leftSubTreeHeight = this.#_getAvlTreeHeight(node.left);
    const rightSubTreeHeight = this.#_getAvlTreeHeight(node.right);

    // 如果左子树高度大于右子树高度，且大不止1，说明不平衡，且这里隐含一个逻辑，就是左子树的左子树一定不为空
    if (leftSubTreeHeight - rightSubTreeHeight > 1) {
        // 左左情况，直接右旋
        if (this.#_getAvlTreeHeight(node.left.left) >= this.#_getAvlTreeHeight(node.left.right)) {
            node = this.#_rotateRight(node)
        } else {
            // 左右情况，先左旋后右旋
            node = this.#_rotateLeftRight(node);
        }
    } else if (rightSubTreeHeight - leftSubTreeHeight > 1) {
        // 右右情况，直接左旋
        if (this.#_getAvlTreeHeight(node.right.right) >= this.#_getAvlTreeHeight(node.right.left)) {
            node = this.#_rotateLeft(node)
        } else {
            // 右左情况，先右旋，后左旋
            node = this.#_rotateRightLeft(node)
        }
    }

    // 返回平衡后的子树的新的根节点
    return node;
}

基本操作

插入

插入新的节点，并返回新的根节点

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
// 这段代码最需要理解的地方是，这个递归函数每次返回的是在以node节点为根节点的树上插入newNode并平衡后的新子树的根节点。
#_insertHelper(node, newNode) {
    // 如果是根节点为空，则直接把newNode当作根节点返回
    if (node === null) {
        return newNode;
    }
    // 如果新节点的值比当前节点小
    if (newNode.val < node.val) {
        // 左节点为空，那么直接就把新节点当作左子节点
        if (node.left === null) {
            node.left = newNode;
        } else {
            // 如果左子节点不为空，那么在左子树中插入新节点，并把插入后返回的新子树的根节点作为左子树
            node.left = this.#_insertHelper(node.left, newNode);
        }
    } else {
        // 右节点为空，那么直接就把新节点当作右子节点
        if (node.right === null) {
            node.right = newNode;
        } else {
            // 如果右子节点不为空，那么在右子树中插入新节点，并把插入后返回的新子树的根节点作为右子树
            node.right = this.#_insertHelper(node.right, newNode);
        }
    }

    // 插入后进行平衡，并返回平衡后的新子树的根节点
    return this.#_balance(node)
}

insert(newNode) {
    // 从根节点开始插入newNode，并且重新赋值根节点，保证下次调用的正确
    this.#_root = this.#_insertHelper(this.#_root, newNode);
}

查找

这个比较简单，和BST没什么区别

1
2
3
4
5
6
7
8
9
10
11
12
13
14
 #_searchHelper(node, val) {
    if (node === null) return null;
    if (node.val === val) {
        return node;
    } else if (val < node.val) {
        return this.#_searchHelper(node.left, val)
    } else {
        return this.#_searchHelper(node.right, val);
    }
}

search(val) {
    return this.#_searchHelper(this.#_root, val);
}

删除

在以node为根节点的子树上删除值为val的节点，并返回平衡后新的树的根节点

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
#_removeHelper(node, val) {
    // 如果node为空，说明树为空，不需要删除，新树的根节点也是null
    if (node === null) return null;
    // 如果找到了要删除的节点
    if (val === node.val) {
        // 要删除的节点左右子树都有
        if (node.left && node.right) {
            // 那就找到该节点右子树
            let p = node.right;
            // 然后在右子树中一直找左子树
            while(p.left) {
                p = p.left;
            }
            // 然后把结果替换到当前节点中，因为右子树中最左的子节点一定小于右子树中其他的节点，同时大雨左子树中所有节点
            node.val = p.val;
            // 上一步只是把节点值替换，但是节点还在，我们需要右子树中删除该节点，然后返回平衡后新的右子树作为node的右子树。
            // 而记住，此时我们要删除的是p.val了，而不是val，因为我们已经把p.val放到了node.val，且这个p有个性质就是它最起码没有左子节点了，座椅下一次会进入紧接着的else逻辑
            node.right = this.#_removeHelper(node.right, p.val)
        } else {
            // 如果不是左右子树都有，直接就替换上来，然后删除原本的就好
            let p = node.left !== null ? node.left : node.right;
            node = null;
            return p;
        }
    } else if (val < node.val) {
        // 如果要删除的值比当前小，就去左子树删除
        node.left = this.#_removeHelper(node.left, val)
    } else {
        // 如果要删除的值比当前大，就去右子树删除
        node.right = this.#_removeHelper(node.right, val);
    }

    // 返回平衡后新的子树根节点
    return this.#_balance(node);
}

remove(val) {
    return this.#_removeHelper(this.#_root, val);
}

完整代码和测试

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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
class AvlTree{
    #_rotateLeft(node) {
        const avlNode = node.right;
        node.right = avlNode.left;
        avlNode.left = node;
        return avlNode;
    }

    #_rotateRight(node) {
        const avlNode = node.left;
        node.left = avlNode.right;
        avlNode.right = node;
        return avlNode;
    }

    #_rotateLeftRight(node) {
        node.left = this.#_rotateLeft(node.left);
        return this.#_rotateRight(node)
    }

    #_rotateRightLeft(node) {
        node.right = this.#_rotateRight(node.right);
        return this.#_rotateLeft(node);
    }

    #_getAvlTreeHeight(node) {
        if (node === null) {
            return 0;
        }

        const leftHeight = this.#_getAvlTreeHeight(node.left);
        const rightHeight = this.#_getAvlTreeHeight(node.right);

        return Math.max(leftHeight, rightHeight) + 1;
    }

    #_balance(node) {
        if (node === null) {
            return node;
        }

        const leftSubTreeHeight = this.#_getAvlTreeHeight(node.left);
        const rightSubTreeHeight = this.#_getAvlTreeHeight(node.right);

        if (leftSubTreeHeight - rightSubTreeHeight > 1) {
            if (this.#_getAvlTreeHeight(node.left.left) >= this.#_getAvlTreeHeight(node.left.right)) {
                node = this.#_rotateRight(node)
            } else {
                node = this.#_rotateLeftRight(node);
            }
        } else if (rightSubTreeHeight - leftSubTreeHeight > 1) {
            if (this.#_getAvlTreeHeight(node.right.right) >= this.#_getAvlTreeHeight(node.right.left)) {
                node = this.#_rotateLeft(node)
            } else {
                node = this.#_rotateRightLeft(node)
            }
        }

        return node;
    }

    #_insertHelper(node, newNode) {
        if (node === null) {
            return newNode;
        }
        if (newNode.val < node.val) {
            if (node.left === null) {
                node.left = newNode;
            } else {
                node.left = this.#_insertHelper(node.left, newNode);
            }
        } else {
            if (node.right === null) {
                node.right = newNode;
            } else {
                node.right = this.#_insertHelper(node.right, newNode);
            }
        }
        return this.#_balance(node)
    }

    insert(newNode) {
        this.#_root = this.#_insertHelper(this.#_root, newNode);
    }

    #_inOrder(node) {
        const stack = [node];
        const result = [];
        while(stack.length) {
            let current = stack.pop();
            if (current !== null) {
                current.right && stack.push(current.right);

                stack.push(current, null);

                current.left && stack.push(current.left);
            } else {
                current = stack.pop();
                result.push(current.val);
            }
        }

        console.log(result);
    } 

    print() {
        if (this.#_root === null) return;
        this.#_inOrder(this.#_root);
        console.log(this.#_getAvlTreeHeight(this.#_root.left))
        console.log(this.#_getAvlTreeHeight(this.#_root.right))
    }

    #_searchHelper(node, val) {
        if (node === null) return null;
        if (node.val === val) {
            return node;
        } else if (val < node.val) {
            return this.#_searchHelper(node.left, val)
        } else {
            return this.#_searchHelper(node.right, val);
        }
    }

    search(val) {
        return this.#_searchHelper(this.#_root, val);
    }

    #_removeHelper(node, val) {
        if (node === null) return null;
        if (val === node.val) {
            if (node.left && node.right) {
                let p = node.right;
                while(p.left) {
                    p = p.left;
                }
                node.val = p.val;
                node.right = this.#_removeHelper(node.right, p.val)
            } else {
                let p = node.left !== null ? node.left : node.right;
                node = null;
                return p;
            }
        } else if (val < node.val) {
            node.left = this.#_removeHelper(node.left, val)
        } else {
            node.right = this.#_removeHelper(node.right, val);
        }

        return this.#_balance(node);
    }

    remove(val) {
        return this.#_removeHelper(this.#_root, val);
    }

    #_root = null;
}

const avlTree = new AvlTree();
avlTree.insert(new TreeNode(0));
avlTree.insert(new TreeNode(2));
avlTree.insert(new TreeNode(3));
avlTree.insert(new TreeNode(4));
avlTree.insert(new TreeNode(6));
avlTree.insert(new TreeNode(8));
avlTree.insert(new TreeNode(9));
avlTree.insert(new TreeNode(14));
avlTree.insert(new TreeNode(15));
avlTree.insert(new TreeNode(16));
avlTree.insert(new TreeNode(17));
avlTree.insert(new TreeNode(18));
avlTree.insert(new TreeNode(19));
avlTree.insert(new TreeNode(20));
avlTree.insert(new TreeNode(21));
avlTree.insert(new TreeNode(11));
avlTree.insert(new TreeNode(31));
avlTree.insert(new TreeNode(61));
avlTree.insert(new TreeNode(111));
avlTree.insert(new TreeNode(221));

avlTree.remove(20);
avlTree.print()

avlTree.remove(1);
avlTree.print()

avlTree.remove(11);
avlTree.print()

avlTree.remove(221);
avlTree.print()

avlTree.remove(8);
avlTree.print()

avlTree.remove(4);
avlTree.print()