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 = 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);