JavaScript Implementation of AVL Tree

Posted on Edited on In Waline:

Recently tried to implement the AVL tree, found that this thing to write up or a lot of details, messed up half a day, so here to record.

An AVL tree is a balanced binary search tree. On the one hand, it is a BST, i.e., a binary search tree, and on the other hand, it is balanced, i.e., the difference between the height of the left and right subtrees of any subtree whose node is the root node is less than one.

Fundamentals

Not much to say about BST, it is relatively simple, that is, the left child of each node is smaller than the root node, the right child is larger than the root node, just pay attention when inserting.

But there is a problem with BST, for example, inserting 1, 2, 3, 4, 5, 6 in sequence, then we will end up with a tree with only right children, in fact, this BST has degenerated into an ordinary chain table, so we need some way to keep this tree balanced, and the way to keep it balanced is as follows:

Specific implementation

Four types of imbalance

Left-right situation

In this case to rotate right and then return the new root node after the rotation, the code is as follows:

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

Right-right situation

In this case to rotate left and then return the new root node after the rotation, the code is as follows:

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

Circumstances

In this case, the left child node of the root node should first be left rotated to become the left-left case, then the root node should be right rotated, and then the new root node should be returned after the rotation, with the following code:

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

Right-left case

In this case, the right child node of the root node should first be rotated right to become the right-right case, then the root node should be rotated left, and then the new root node should be returned after the rotation, with the following code:

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

Fix the imbalance

Get the height of the node

When we do balancing, we need to determine whether we need to balance and what kind of imbalance is the case so that we can choose a different rotation method

1
2
3
4
5
6
7
8
9
10
11
12
13
#_getAvlTreeHeight(node) {
    // The height of the empty node is 0
    if (node= null) {
        return 0;
    }

    // Post-order traversal is used here, because the result of each recursion requires the result of the child node recursion
    const leftHeight = this.#_getAvlTreeHeight(node.left);
    const rightHeight = this.#_getAvlTreeHeight(node.right);

    // The node height is the greater of the left and right subtree heights plus 1
    return Math.max(leftHeight, rightHeight) + 1;
}

Balance

Balances the subtree with node as the root node and returns the new root 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) {
    // Empty nodes do not need to be balanced, return empty nodes directly
    if (node= null) {
        return node;
    }

    // Get the height of the left and right subtrees
    const leftSubTreeHeight = this.#_getAvlTreeHeight(node.left);
    const rightSubTreeHeight = this.#_getAvlTreeHeight(node.right);

    // If the height of the left subtree is greater than the height of the right subtree by more than 1, it is unbalanced and there is an implied logic that the left subtree of the left subtree must not be empty
    if (leftSubTreeHeight - rightSubTreeHeight > 1) {
        // Left-left situation, direct right rotation
        if (this.#_getAvlTreeHeight(node.left.left) >= this.#_getAvlTreeHeight(node.left.right)) {
            node = this.#_rotateRight(node)
        } else {
            // Left and right situation, first left rotation then right rotation
            node = this.#_rotateLeftRight(node);
        }
    } else if (rightSubTreeHeight - leftSubTreeHeight > 1) {
        // Right-right situation, direct left rotation
        if (this.#_getAvlTreeHeight(node.right.right) >= this.#_getAvlTreeHeight(node.right.left)) {
            node = this.#_rotateLeft(node)
        } else {
            // Right-left situation, first right, then left
            node = this.#_rotateRightLeft(node)
        }
    }

    // Return the new root node of the balanced subtree
    return node;
}

Basic operation

Insert

Inserts a new node and returns the new root 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
// The most important thing to understand about this code is that this recursive function returns the root of a new subtree each time a newNode is inserted and balanced on a tree with a node node as the root node.
#_insertHelper(node, newNode) {
    // If the root node is empty, return the newNode directly as the root node
    if (node= null) {
        return newNode;
    }
    // If the value of the new node is smaller than the current node
    if (newNode.val < node.val) {
        // If the left node is empty, then the new node will be treated as the left child node directly
        if (node.left= null) {
            node.left = newNode;
        } else {
            // If the left child node is not empty, then insert a new node into the left subtree and use the root node of the new subtree returned after the insertion as the left subtree
            node.left = this.#_insertHelper(node.left, newNode);
        }
    } else {
        // If the right node is empty, then the new node is treated as the right child node directly
        if (node.right= null) {
            node.right = newNode;
        } else {
            // If the right child node is not empty, then insert a new node into the right subtree and use the root node of the new subtree returned after the insertion as the right subtree
            node.right = this.#_insertHelper(node.right, newNode);
        }
    }

    // Balance after insertion and return the root node of the new subtree after balancing
    return this.#_balance(node)
}

insert(newNode) {
    // Insert the newNode from the root, and reassign the root to ensure the next call is correct
    this.#_root = this.#_insertHelper(this.#_root, newNode);
}

Find

This is relatively simple, and not much different from 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);
}

Delete

Delete the node with the value val from the subtree with node as the root and return the root of the new tree after balancing

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) {
    // If node is null, the tree is empty and does not need to be deleted, and the root node of the new tree is also null
    if (node= null) return null;
    // If the node to be deleted is found
    if (trap= node.val) {
        // The node to be deleted has both left and right subtrees
        if (node.left && node.right) {
            // Then find the right subtree of the node
            let p = node.right;
            // Then look for the left subtree all the way down the right subtree
            while(p.left) {
                p = p.left;
            }
            // Then replace the result in the current node, because the leftmost child node in the right subtree must be smaller than the rest of the nodes in the right subtree, while raining all the nodes in the left subtree
            node.val = p.val;
            // The previous step just replaces the node value, but the node is still there. We need to remove the node from the right subtree and return the new right subtree as the right subtree of the node after balancing.
            // And remember, at this point we are deleting p.val, not val, because we have already put p.val into node.val, and this p has the property that it has no left child node at least, and the next time the seat will enter the immediately following else logic
            node.right = this.#_removeHelper(node.right, p.val)
        } else {
            // If you don't have both left and right subtrees, just replace them and delete the original ones
            let p = node.left ! null ? node.left : node.right;
            node = null;
            return p;
        }
    } else if (val < node.val) {
        // If the value to be deleted is smaller than the current one, go to the left subtree and delete it
        node.left = this.#_removeHelper(node.left, val)
    } else {
        // If the value to be deleted is larger than the current one, go to the right subtree and delete it
        node.right = this.#_removeHelper(node.right, val);
    }

    // Return the root of the new subtree after balancing
    return this.#_balance(node);
}

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

Complete code and tests

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 (trap= 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()