diff --git a/src/data-structures/tree/BinaryTreeNode.js b/src/data-structures/tree/BinaryTreeNode.js index 5498b9d6..cca840da 100644 --- a/src/data-structures/tree/BinaryTreeNode.js +++ b/src/data-structures/tree/BinaryTreeNode.js @@ -32,13 +32,17 @@ export default class BinaryTreeNode { setLeft(node) { this.left = node; - this.left.parent = this; + if (node) { + this.left.parent = this; + } return this; } setRight(node) { this.right = node; - this.right.parent = this; + if (node) { + this.right.parent = this; + } return this; } diff --git a/src/data-structures/tree/__test__/BinaryTreeNode.test.js b/src/data-structures/tree/__test__/BinaryTreeNode.test.js index c32123a1..9160a965 100644 --- a/src/data-structures/tree/__test__/BinaryTreeNode.test.js +++ b/src/data-structures/tree/__test__/BinaryTreeNode.test.js @@ -155,4 +155,22 @@ describe('BinaryTreeNode', () => { expect(right.height).toBe(0); expect(root.balanceFactor).toBe(-1); }); + + it('should set null for left and right node', () => { + const root = new BinaryTreeNode(2); + const left = new BinaryTreeNode(1); + const right = new BinaryTreeNode(3); + + root.setLeft(left); + root.setRight(right); + + expect(root.left.value).toBe(1); + expect(root.right.value).toBe(3); + + root.setLeft(null); + root.setRight(null); + + expect(root.left).toBeNull(); + expect(root.right).toBeNull(); + }); });