From 47fcfabae7c1b94fb6293cecce5a65baf5679935 Mon Sep 17 00:00:00 2001
From: Revone <zrwusa@gmail.com>
Date: Wed, 18 Oct 2023 18:28:36 +0800
Subject: [PATCH] [binarytree] modified the getDepth method to adhere to the
 proper definition of "tree depth."

---
 src/data-structures/binary-tree/binary-tree.ts     | 14 +++++++++-----
 .../binary-tree/binary-tree.test.ts                | 13 +++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/data-structures/binary-tree/binary-tree.ts b/src/data-structures/binary-tree/binary-tree.ts
index 98f058e..4f1e942 100644
--- a/src/data-structures/binary-tree/binary-tree.ts
+++ b/src/data-structures/binary-tree/binary-tree.ts
@@ -350,16 +350,20 @@ export class BinaryTree<N extends BinaryTreeNode<N['val'], N> = BinaryTreeNode>
 
   /**
    * The function calculates the depth of a node in a binary tree.
-   * @param {N | BinaryTreeNodeKey | null} beginRoot - The `beginRoot` parameter can be one of the following:
+   * @param {N | BinaryTreeNodeKey | null} distNode - The `distNode` parameter can be any node of the tree
+   * @param {N | BinaryTreeNodeKey | null} beginRoot - The `beginRoot` parameter can be the predecessor node of distNode
    * @returns the depth of the given node or binary tree.
    */
-  getDepth(beginRoot: N | BinaryTreeNodeKey | null = this.root): number {
+  getDepth(distNode: N | BinaryTreeNodeKey | null, beginRoot: N | BinaryTreeNodeKey | null = this.root): number {
+    if (typeof distNode === 'number') distNode = this.get(distNode, 'key');
     if (typeof beginRoot === 'number') beginRoot = this.get(beginRoot, 'key');
-
     let depth = 0;
-    while (beginRoot?.parent) {
+    while (distNode?.parent) {
+      if (distNode === beginRoot) {
+        return depth;
+      }
       depth++;
-      beginRoot = beginRoot.parent;
+      distNode = distNode.parent;
     }
     return depth;
   }
diff --git a/test/unit/data-structures/binary-tree/binary-tree.test.ts b/test/unit/data-structures/binary-tree/binary-tree.test.ts
index abef0e9..ab16be2 100644
--- a/test/unit/data-structures/binary-tree/binary-tree.test.ts
+++ b/test/unit/data-structures/binary-tree/binary-tree.test.ts
@@ -105,6 +105,19 @@ describe('BinaryTree', () => {
     expect(binaryTree.has(4)).toBe(false);
   });
 
+  test('should getDepth return correct depth', () => {
+    binaryTree.add(1);
+    expect(binaryTree.getDepth(1)).toBe(0);
+    binaryTree.add(2);
+    expect(binaryTree.getDepth(2)).toBe(1);
+    binaryTree.add(3);
+    expect(binaryTree.getDepth(3, 1)).toBe(1);
+    binaryTree.add(4);
+    expect(binaryTree.getDepth(4, 1)).toBe(2);
+    expect(binaryTree.getDepth(4)).toBe(2);
+    expect(binaryTree.getDepth(4, 2)).toBe(1);
+  });
+
   test('should traverse in-order', () => {
     binaryTree.add(4);
     binaryTree.add(2);