Skip to content

Commit 04e533e

Browse files
Bladefidztrekhleb
authored andcommitted
Add remove method (trekhleb#33)
Remove node in AvlTree with auto balancing. Fix issue: trekhleb#13
1 parent ada4537 commit 04e533e

File tree

2 files changed

+80
-0
lines changed

2 files changed

+80
-0
lines changed

src/data-structures/tree/avl-tree/AvlTree.js

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,48 @@ export default class AvlTree extends BinarySearchTree {
1616
}
1717
}
1818

19+
remove(value) {
20+
const nodeToRemove = this.root.find(value);
21+
22+
if (!nodeToRemove) {
23+
throw new Error('Item not found in the tree');
24+
}
25+
26+
// Recursively find target node, if found then delete and balance.
27+
// return nodeToRemove.value;
28+
this.root = this.removeRecv(this.root, nodeToRemove);
29+
}
30+
31+
removeRecv(origin, node) {
32+
let newOrigin = origin;
33+
if (origin.value > node.value) {
34+
// Recursively traversing from left
35+
newOrigin.left = this.removeRecv(origin.left, node);
36+
} else if (origin.value < node.value) {
37+
// Recursively traversing from right
38+
newOrigin.right = this.removeRecv(origin.right, node);
39+
} else {
40+
if (origin.left == null) {
41+
// Forget right node
42+
return origin.right;
43+
}
44+
if (origin.right == null) {
45+
// Forget left node
46+
return origin.left;
47+
}
48+
49+
// Recursively find min node from left subtree
50+
// more efficient traversing
51+
const parent = Object.assign({}, origin);
52+
newOrigin = parent.right.findMin();
53+
newOrigin.right = this.deleteMin(parent.right);
54+
newOrigin.left = parent.left;
55+
}
56+
57+
// Balance and return root node
58+
return this.balance(newOrigin);
59+
}
60+
1961
/**
2062
* @param {*} value
2163
* @return {boolean}
@@ -48,6 +90,8 @@ export default class AvlTree extends BinarySearchTree {
4890
this.rotateRightLeft(node);
4991
}
5092
}
93+
// Return the heap to avoid referenceError
94+
return node;
5195
}
5296

5397
/**
@@ -156,4 +200,15 @@ export default class AvlTree extends BinarySearchTree {
156200
// Attach rootNode to the left of rightNode.
157201
rightNode.setLeft(rootNode);
158202
}
203+
204+
deleteMin(node) {
205+
// Forget right node if has value
206+
if (node.left == null) return node.right;
207+
208+
const newNode = node;
209+
newNode.left = this.deleteMin(node.left);
210+
211+
// Balance and return root node
212+
return this.balance(newNode);
213+
}
159214
}

src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,31 @@ describe('AvlTree', () => {
186186
expect(tree.toString()).toBe('6,8,9,18,21,22,43');
187187
});
188188

189+
it('should keep balance after removal', () => {
190+
const tree = new AvlTree();
191+
192+
tree.insert(1);
193+
tree.insert(2);
194+
tree.insert(3);
195+
tree.insert(4);
196+
tree.insert(5);
197+
tree.insert(6);
198+
tree.insert(7);
199+
tree.insert(8);
200+
tree.insert(9);
201+
202+
expect(tree.toString()).toBe('1,2,3,4,5,6,7,8,9');
203+
expect(tree.root.height).toBe(3);
204+
205+
tree.remove(8);
206+
tree.remove(9);
207+
208+
expect(tree.contains(8)).toBeFalsy();
209+
expect(tree.contains(9)).toBeFalsy();
210+
expect(tree.toString()).toBe('1,2,3,4,5,6,7');
211+
expect(tree.root.height).toBe(2);
212+
});
213+
189214
it('should do left right rotation and keeping left right node safe', () => {
190215
const tree = new AvlTree();
191216

0 commit comments

Comments
 (0)