diff --git a/src/data-structures/trees/binary-search-tree.js b/src/data-structures/trees/binary-search-tree.js index 5e9ffb8..e7a5cc4 100644 --- a/src/data-structures/trees/binary-search-tree.js +++ b/src/data-structures/trees/binary-search-tree.js @@ -19,23 +19,24 @@ class BinarySearchTree { * @returns {BinaryTreeNode} newly added node */ add(value) { - const newNode = new BinaryTreeNode(value); + let node = new BinaryTreeNode(value); if (this.root) { const { found, parent } = this.findNodeAndParent(value); // <1> if (found) { // duplicated: value already exist on the tree found.meta.multiplicity = (found.meta.multiplicity || 1) + 1; // <2> + node = found; } else if (value < parent.value) { - parent.setLeftAndUpdateParent(newNode); + parent.setLeftAndUpdateParent(node); } else { - parent.setRightAndUpdateParent(newNode); + parent.setRightAndUpdateParent(node); } } else { - this.root = newNode; + this.root = node; } this.size += 1; - return newNode; + return node; } // end::add[] diff --git a/src/data-structures/trees/binary-search-tree.spec.js b/src/data-structures/trees/binary-search-tree.spec.js index 854e31f..22508eb 100644 --- a/src/data-structures/trees/binary-search-tree.spec.js +++ b/src/data-structures/trees/binary-search-tree.spec.js @@ -66,7 +66,7 @@ describe('Binary Search Tree', () => { it('should deal with duplicates', () => { const root = bst.add(1); expect(root.meta.multiplicity).toBe(undefined); - bst.add(1); + expect(bst.add(1)).toBe(root); // should return existing expect(bst.size).toBe(2); expect(root.toValues()).toMatchObject({ value: 1, parent: null, left: null, right: null, @@ -262,7 +262,7 @@ describe('Binary Search Tree', () => { }); it('should remove duplicates', () => { - bst.add(40); // add duplicate + expect(bst.add(40)).toBe(n40); // add duplicate expect(n40.meta.multiplicity).toBe(2); expect(bst.remove(40)).toBe(true);