Coding Interview Questions | Diameter of Binary Tree

January 2nd, 2024

Diameter of a Binary Tree Problem Introduction:

This summary will talk about my solution to the diameter of a binary tree problem as seen on leetcode here.

How do you solve the Diameter of a Binary Tree problem?:

Solution

const getHeight = (
  node: TreeNode | null,
  nodeHeightMap: Map<TreeNode, number>
) => {
  if (!node) {
    return 0;
  }
  if (nodeHeightMap.has(node)) {
    return nodeHeightMap.get(node);
  }

  const leftHeight = getHeight(node.left, nodeHeightMap);
  const rightHeight = getHeight(node.right, nodeHeightMap);

  const height = Math.max(leftHeight, rightHeight) + 1;

  if (!nodeHeightMap.has(node)) {
    nodeHeightMap.set(node, height);
  }

  return height;
};

function diameterOfBinaryTree(root: TreeNode | null): number {
  if (!root) {
    return 0;
  }

  const nodeStack = [root];
  let maxHeight = 0;
  const nodeHeightMap = new Map<TreeNode, number>();

  while (nodeStack.length > 0) {
    const currentNode = nodeStack.pop()!;

    const currentDiameter =
      getHeight(currentNode.left, nodeHeightMap) +
      getHeight(currentNode.right, nodeHeightMap);
    if (currentDiameter > maxHeight) {
      maxHeight = currentDiameter;
    }

    if (currentNode.left) {
      nodeStack.push(currentNode.left);
    }

    if (currentNode.right) {
      nodeStack.push(currentNode.right);
    }
  }

  return maxHeight;
}

Diameter of a Binary Tree Solution Summary:

getHeight Function - Recursive Height Calculation with Memoization

1. Base Case: First we define our base case. If the node is null, return 0.
2. Memoization Check: Then, if the height of the node is already calculated, return the stored value.
3. Recursive Calls: Next we calculate the height of left and right subtrees using recursion.
4. Memoization: Store the calculated height in the nodeHeightMap for future reference.
5. Return: Finally, we return the calculated height to the diameterOfBinaryTree function.

diameterOfBinaryTree Function - Iterative Diameter Calculation Using Stack and Memoization

1. Base Case: Again, we define our base case. If the root is null, return 0 (empty tree).
2. Initialization: We then initialize a stack (nodeStack), maximum diameter (maxHeight), and a memoization map (nodeHeightMap).
3. Main Loop: Then, we perform a stack-based iterative traversal of the tree.
   • First we pop a node from the stack.
   • Next, we calculate the diameter of the current node using the getHeight function.
   • We then use that calculated height and update maxHeight if the current diameter is greater.
   • Finally, we push left and right child nodes onto the stack if they exist.

Complexities

• Our overall time complexity is O(n) as we will iterate through all nodes in the tree once.
• Our worst case space complexity is also O(n) given the map employed in the getHeight function.