From 857edbf3a8bdb063a7cf1ebe07e4d5f37c765ae3 Mon Sep 17 00:00:00 2001
From: Oleksii Trekhleb <trehleb@gmail.com>
Date: Wed, 4 Apr 2018 09:11:26 +0300
Subject: [PATCH] Add possibility for tree nodes to have height.

---
 src/data-structures/tree/BinaryTreeNode.js    | 11 ++++++
 .../tree/__test__/BinaryTreeNode.test.js      | 38 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/src/data-structures/tree/BinaryTreeNode.js b/src/data-structures/tree/BinaryTreeNode.js
index 9db31ea3..e0ed5668 100644
--- a/src/data-structures/tree/BinaryTreeNode.js
+++ b/src/data-structures/tree/BinaryTreeNode.js
@@ -6,6 +6,17 @@ export default class BinaryTreeNode {
     this.value = value;
   }
 
+  get height() {
+    if (!this.left && !this.left) {
+      return 0;
+    }
+
+    const leftHeight = this.left ? this.left.height : 0;
+    const rightHeight = this.right ? this.right.height : 0;
+
+    return Math.max(leftHeight, rightHeight) + 1;
+  }
+
   setLeft(node) {
     this.left = node;
     this.left.parent = this;
diff --git a/src/data-structures/tree/__test__/BinaryTreeNode.test.js b/src/data-structures/tree/__test__/BinaryTreeNode.test.js
index 54535e53..ffce779f 100644
--- a/src/data-structures/tree/__test__/BinaryTreeNode.test.js
+++ b/src/data-structures/tree/__test__/BinaryTreeNode.test.js
@@ -100,6 +100,44 @@ describe('BinaryTreeNode', () => {
     expect(rootNode.replaceChild(rootNode.right, replacementNode)).toBeTruthy();
     expect(rootNode.traverseInOrder()).toEqual([1, 2, 5]);
 
+    expect(rootNode.replaceChild(rootNode.left, replacementNode)).toBeTruthy();
+    expect(rootNode.traverseInOrder()).toEqual([5, 2, 5]);
+
     expect(rootNode.replaceChild(new BinaryTreeNode(), new BinaryTreeNode())).toBeFalsy();
   });
+
+  it('should calculate node height', () => {
+    const root = new BinaryTreeNode(1);
+    const left = new BinaryTreeNode(3);
+    const right = new BinaryTreeNode(2);
+    const grandLeft = new BinaryTreeNode(5);
+    const grandRight = new BinaryTreeNode(6);
+    const grandGrandLeft = new BinaryTreeNode(7);
+
+    expect(root.height).toBe(0);
+
+    root
+      .setLeft(left)
+      .setRight(right);
+
+    expect(root.height).toBe(1);
+    expect(left.height).toBe(0);
+
+    left
+      .setLeft(grandLeft)
+      .setRight(grandRight);
+
+    expect(root.height).toBe(2);
+    expect(left.height).toBe(1);
+    expect(grandLeft.height).toBe(0);
+    expect(grandRight.height).toBe(0);
+
+    grandLeft.setLeft(grandGrandLeft);
+
+    expect(root.height).toBe(3);
+    expect(left.height).toBe(2);
+    expect(grandLeft.height).toBe(1);
+    expect(grandRight.height).toBe(0);
+    expect(grandGrandLeft.height).toBe(0);
+  });
 });