Problem Description

You are given the root of a binary tree. Your task is to replace the value of each node with the sum of all its cousin nodes' values.

What are cousin nodes?

• Two nodes are cousins if they are at the same depth (level) in the tree but have different parents
• The depth of a node is the number of edges from the root to that node

Key Points:

• Sibling nodes (nodes with the same parent) are NOT cousins to each other
• Nodes at the same level but with different parents ARE cousins
• After replacement, each node's value should equal the sum of all its cousins' values

Example Understanding: If a tree has nodes at depth 2 with values [4, 5, 6, 7] where:

• Nodes 4 and 5 share parent A
• Nodes 6 and 7 share parent B

Then:

• Node 4 has cousins 6 and 7, so it gets replaced with 6 + 7 = 13
• Node 5 has cousins 6 and 7, so it gets replaced with 6 + 7 = 13
• Node 6 has cousins 4 and 5, so it gets replaced with 4 + 5 = 9
• Node 7 has cousins 4 and 5, so it gets replaced with 4 + 5 = 9

The function should return the root of the modified tree where all nodes have been updated with their cousin sum values.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: A binary tree is a special type of graph (specifically, a connected acyclic graph where each node has at most two children).

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree structure where we need to traverse nodes and process them based on their relationships (parent-child and cousin relationships).

DFS

• Yes: We arrive at DFS as the recommended approach.

Conclusion: The flowchart suggests using DFS (Depth-First Search) for this problem.

Why DFS is Appropriate

The DFS pattern fits perfectly because:

1. Level-wise Information Gathering: We need to traverse the entire tree to collect sum information at each depth level. DFS allows us to systematically visit all nodes while tracking their depth.

2. Parent-Child Relationship Tracking: To identify cousins (same depth, different parents), we need to traverse the tree while maintaining parent context. DFS naturally maintains this context through its recursive call stack.

3. Two-Pass Strategy:

   • First DFS pass: Calculate the sum of all nodes at each depth level
   • Second DFS pass: Update each node's value based on cousin relationships (total sum at that level minus sibling values)

4. Tree Structure Navigation: DFS is ideal for binary tree problems where we need to process nodes based on their structural relationships rather than just their values.

Intuition

The key insight is recognizing that cousins are all nodes at the same level except for siblings. This means for any node, its cousin sum equals the total sum of all nodes at that level minus the sum of itself and its sibling.

Let's think through this step by step:

1. What information do we need?

   • To calculate cousin sums, we need to know the total sum at each depth level
   • We also need to know which nodes are siblings (share the same parent)

2. Breaking down the cousin relationship:

   • If nodes A and B are at depth d and have the same parent, they are siblings
   • All other nodes at depth d are their cousins
   • So: cousin_sum = total_sum_at_depth_d - (A.val + B.val)

3. Why two passes?

   • First pass: We can't calculate cousin sums without knowing the total at each level first. So we traverse the tree once to build an array s where s[depth] stores the sum of all nodes at that depth.
   • Second pass: Now that we know level sums, we can update each node. But here's the clever part - instead of updating a node with its own cousin sum, we update its children with their cousin sums. Why? Because when we're at a parent node, we can easily calculate the sum of its children (the siblings), and therefore determine what their cousin sum should be.

4. The parent's perspective trick:

   • When we're at a parent node during the second DFS, we know:
     • The total sum at the children's level (from our first pass)
     • The values of both children (direct access)
   • This allows us to calculate: cousin_sum = level_sum - sibling_sum
   • We then assign this value to each child

5. Edge cases handled naturally:

   • The root has no cousins (it's alone at depth 0), so we set it to 0
   • If a node has only one child, we still calculate correctly since a missing child contributes 0 to the sibling sum
   • Leaf nodes have no children to update, so the recursion naturally terminates

This approach elegantly solves the problem by recognizing that the cousin relationship is best computed from the parent's perspective, where we have easy access to sibling information.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Approach

The solution implements a two-pass DFS strategy as outlined in the reference approach:

First Pass - Building Level Sums

The first DFS traversal (dfs1) collects the sum of node values at each depth:

def dfs1(root: Optional[TreeNode], depth: int):
    if root is None:
        return
    if len(s) <= depth:
        s.append(0)
    s[depth] += root.val
    dfs1(root.left, depth + 1)
    dfs1(root.right, depth + 1)

Key components:

• Data Structure: A list s where s[depth] stores the sum of all node values at that depth
• Dynamic list growth: If we encounter a new depth level, we extend the list
• Recursive traversal: Visit each node once, adding its value to the appropriate depth sum
• Time complexity: O(n) where n is the number of nodes

Second Pass - Updating Node Values

The second DFS traversal (dfs2) updates each node's children with their cousin sums:

def dfs2(root: Optional[TreeNode], depth: int):
    sub = (root.left.val if root.left else 0) + (
        root.right.val if root.right else 0
    )
    depth += 1
    if root.left:
        root.left.val = s[depth] - sub
        dfs2(root.left, depth)
    if root.right:
        root.right.val = s[depth] - sub
        dfs2(root.right, depth)

The algorithm:

1. Calculate sub = sum of current node's children (siblings to each other)
2. For each child that exists:
   • Set its value to s[depth] - sub (total at that level minus sibling sum)
   • Recursively process that child's subtree

Why this works:

• When at a parent node, we know both children (if they exist)
• These children are siblings to each other
• Their cousin sum = (total sum at their level) - (sum of siblings)
• Since we're modifying values as we go, we must calculate the sibling sum before updating

Main Function Flow

s = []
dfs1(root, 0)          # Build level sums
root.val = 0           # Root has no cousins
dfs2(root, 0)          # Update all other nodes
return root

Special handling:

• The root node is set to 0 since it has no cousins (it's the only node at depth 0)
• The dfs2 function starts from the root but only updates its children and below

Complexity Analysis:

• Time Complexity: O(n) - we visit each node exactly twice
• Space Complexity: O(h) where h is the height of the tree (for recursion stack) plus O(d) for the level sum array where d is the depth of the tree

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Consider this binary tree:

       5
      / \
     4   9
    / \
   1   10

Step 1: First DFS Pass - Calculate Level Sums

Starting from the root, we traverse the tree and build our level sums array s:

• Depth 0: Visit node 5 → s[0] = 5
• Depth 1: Visit nodes 4 and 9 → s[1] = 4 + 9 = 13
• Depth 2: Visit nodes 1 and 10 → s[2] = 1 + 10 = 11

After first pass: s = [5, 13, 11]

Step 2: Update Root

The root has no cousins (it's alone at depth 0), so we set root.val = 0.

Tree now looks like:

       0
      / \
     4   9
    / \
   1   10

Step 3: Second DFS Pass - Update with Cousin Sums

Starting from the root, we update each node's children:

At root (depth 0):

• Calculate sibling sum: sub = 4 + 9 = 13
• Update left child (4): s[1] - sub = 13 - 13 = 0
• Update right child (9): s[1] - sub = 13 - 13 = 0

Tree becomes:

       0
      / \
     0   0
    / \
   1   10

At node 4 (now 0) at depth 1:

• Calculate sibling sum: sub = 1 + 10 = 11
• Update left child (1): s[2] - sub = 11 - 11 = 0
• Update right child (10): s[2] - sub = 11 - 11 = 0

At node 9 (now 0) at depth 1:

• No children, so nothing to update

Final tree:

       0
      / \
     0   0
    / \
   0   0

Verification:

• Node at depth 0: Has no cousins → value = 0 ✓
• Nodes at depth 1 (originally 4 and 9): They are siblings, not cousins to each other → cousin sum = 0 ✓
• Nodes at depth 2 (originally 1 and 10): They are siblings, not cousins to each other → cousin sum = 0 ✓

This example shows how the algorithm correctly identifies that sibling nodes have no cousins at their level when they are the only nodes at that depth with different parents.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm consists of two depth-first search (DFS) traversals:

• dfs1: Traverses the entire tree once to calculate the sum of node values at each depth level. Each node is visited exactly once, resulting in O(n) time.
• dfs2: Traverses the entire tree once more to replace each node's value with the sum of its cousins (nodes at the same depth but with different parents). Each node is visited exactly once, resulting in O(n) time.

Since we perform two sequential traversals of the tree, the total time complexity is O(n) + O(n) = O(n).

Space Complexity: O(n)

The space complexity comes from multiple sources:

• Recursion call stack: In the worst case (a skewed tree), the recursion depth can be O(n), requiring O(n) space for the call stack.
• Array s: This array stores the sum of values at each depth level. In the worst case (a skewed tree), there can be n depth levels, requiring O(n) space.
• Input tree structure: The tree itself occupies O(n) space, but since we're modifying the existing tree in-place rather than creating a new one, this doesn't add to the auxiliary space complexity.

The dominant factor is the maximum between the recursion stack and the array storage, both of which are O(n) in the worst case. Therefore, the overall space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Modifying Values Before Reading Siblings

The Pitfall: A common mistake is updating node values during traversal without first saving the original sibling values. Since we need to calculate sibling_sum to determine cousin sums, if we modify a node's value before reading its sibling, the calculation becomes incorrect.

Incorrect Approach:

def update_node_values(parent: Optional[TreeNode], depth: int) -> None:
    child_depth = depth + 1
  
    # WRONG: Updating left child before reading right child's value
    if parent.left:
        parent.left.val = level_sums[child_depth] - parent.left.val
        update_node_values(parent.left, child_depth)
  
    # Now the sibling sum calculation would be wrong!
    if parent.right:
        # This uses the already modified left.val, not the original
        sibling_sum = parent.left.val + parent.right.val  
        parent.right.val = level_sums[child_depth] - sibling_sum

The Solution: Always calculate and store the sibling sum BEFORE modifying any values:

# Calculate sibling sum first with original values
sibling_sum = (parent.left.val if parent.left else 0) + \
             (parent.right.val if parent.right else 0)

# Then update both children using the same sibling_sum
if parent.left:
    parent.left.val = level_sums[child_depth] - sibling_sum
if parent.right:
    parent.right.val = level_sums[child_depth] - sibling_sum

2. Attempting Bottom-Up Value Updates

The Pitfall: Trying to update a node's value when visiting that node directly, rather than updating it through its parent. This approach fails because you can't determine siblings without access to the parent.

Incorrect Approach:

def update_values(node: Optional[TreeNode], depth: int, parent: Optional[TreeNode]):
    if node is None:
        return
  
    # WRONG: How do we find siblings without parent reference?
    # We'd need to pass extra information about siblings
    node.val = level_sums[depth] - get_sibling_sum(node, parent)  # Complex!
  
    update_values(node.left, depth + 1, node)
    update_values(node.right, depth + 1, node)

The Solution: Update children from their parent node, where both siblings are directly accessible:

def update_node_values(parent: Optional[TreeNode], depth: int):
    # Parent knows both children - can easily calculate sibling sum
    sibling_sum = (parent.left.val if parent.left else 0) + \
                 (parent.right.val if parent.right else 0)
  
    # Update children from parent's perspective
    if parent.left:
        parent.left.val = level_sums[depth + 1] - sibling_sum
    if parent.right:
        parent.right.val = level_sums[depth + 1] - sibling_sum

3. Forgetting Edge Cases

The Pitfall: Not handling special cases properly:

• Root node (has no cousins)
• Single child nodes (only one sibling exists)
• Leaf nodes (have no children to update)

Incorrect Handling:

# WRONG: Forgetting root node special case
def replaceValueInTree(root):
    calculate_level_sums(root, 0)
    update_node_values(root, 0)  # Root value never gets set to 0!
    return root

# WRONG: Not handling single child case
sibling_sum = parent.left.val + parent.right.val  # NullPointerError if one is None!

The Solution: Handle all edge cases explicitly:

# Root has no cousins
root.val = 0

# Safe sibling sum calculation
sibling_sum = (parent.left.val if parent.left else 0) + \
             (parent.right.val if parent.right else 0)

# Leaf nodes naturally handled - the function just returns without updating anything