All Kinds Of Node Depths

All Kinds Of Node Depths

Category

Binary Trees

Difficulty

Very Hard

Problem Statement

Given a binary tree, compute the sum of all node depths for every node treated as the root of its own subtree. For each node in the tree, calculate the sum of depths of all nodes in its subtree (where the node itself has depth 0). Return the total sum of all these per-node depth sums. This is sometimes called the "sum of all subtree node depth sums."

Intuition

The standard "node depths" problem computes the sum of depths from the root. This problem extends it to every subtree. A brute-force approach would compute the depth sum for each node's subtree independently, leading to O(n^2) time. However, we can derive a recursive relationship: when we move from a parent to its child, the child's subtree depth sum can be computed from the parent's using the size of the child's subtree. Additionally, each node contributes to the depth sum of every ancestor — and the contribution to an ancestor at distance d is exactly d. So the total answer is the sum over all pairs (ancestor, descendant) of the distance between them.

Approach

Approach 1: Recursive with subtree info

  1. At each node, compute both the "node depth sum" (sum of depths within its subtree) and the "subtree size."
  2. For a node with left child L and right child R:
    • size(node) = 1 + size(L) + size(R)
    • depthSum(node) = depthSum(L) + depthSum(R) + size(L) + size(R) (each node in the left/right subtree is one level deeper relative to the current node)
  3. The answer is the sum of depthSum(node) for all nodes.
  4. Compute everything in one DFS pass by accumulating the total sum as a global variable.

Approach 2: Contribution counting

  1. Each node at depth d from the root contributes 1 to the depth sum of each of its d ancestors (at distances 1, 2, ..., d).
  2. Equivalently, the total is the sum of d * (number of nodes at depth d in each subtree) for all subtrees, which simplifies to summing f(node) for all nodes, where f(node) counts how many (ancestor, descendant) pairs include this node as the descendant. That count is the node's depth from the root... but we need it for every subtree root, not just the main root.
  3. A cleaner formulation: the total answer equals the sum over all nodes v of depth(v) * (number of ancestors of v, including itself as subtree root). This reduces to summing depth(v) * (depth(v) + 1) / 2 over all nodes, but this overcounts. The correct recursive formulation from Approach 1 is simpler.

Pseudocode

function allKindsOfNodeDepths(root):
    totalSum = 0

    function dfs(node):
        nonlocal totalSum
        if node is null:
            return (0, 0)  # (depthSum, size)

        leftDepthSum, leftSize = dfs(node.left)
        rightDepthSum, rightSize = dfs(node.right)

        currentDepthSum = leftDepthSum + rightDepthSum + leftSize + rightSize
        currentSize = 1 + leftSize + rightSize

        totalSum += currentDepthSum
        return (currentDepthSum, currentSize)

    dfs(root)
    return totalSum

Time & Space Complexity

  • Time: O(n) where n is the number of nodes. Each node is visited exactly once in the DFS.
  • Space: O(h) where h is the height of the tree, for the recursion stack. In the worst case (skewed tree), this is O(n). For a balanced tree, it is O(log n).

Key Insights

  • The depth sum of a subtree rooted at a node can be computed from its children's depth sums and subtree sizes, avoiding redundant computation.
  • Adding size(child) when computing the parent's depth sum accounts for the fact that every node in the child's subtree is one level deeper relative to the parent.
  • This pattern — augmenting DFS return values with extra information (here, subtree size) — is a powerful technique for many binary tree problems.
  • The O(n) solution is a significant improvement over the naive O(n^2) approach of running a depth-sum computation from every node.

Edge Cases

  • Empty tree (null root): return 0.
  • Single node: return 0 (no descendants, depth sum is 0).
  • Perfectly balanced tree: the recursion is well-balanced, and the result is symmetric.
  • Completely skewed tree (linked list): each node's subtree is a prefix of the list; the total grows as O(n^3 / 6).
  • Tree with only left children or only right children: same as skewed case.