Dynamic Programming on Trees

1. Introduction: Escaping the Array

Most developers associate Dynamic Programming strictly with 1D Arrays or 2D Matrices. However, DP is fundamentally about breaking a problem down into overlapping subproblems. What data structure inherently breaks down into identical, smaller subproblems? A Tree!

Every node in a tree acts as the root of its own smaller sub-tree. DP on Trees (also known as Tree DP) involves using Post-Order Traversal (DFS) to solve a problem for the left child and right child, and then using those two solutions to compute the optimal answer for the parent node.

Instead of caching results in a massive 2D array or HashMap (which is slow and cumbersome for tree structures), Tree DP almost exclusively uses Return States. The recursive function returns a small array or custom object containing the DP states for that specific node.

2. Core Concept: The Post-Order State Machine

In array DP, we look backward to dp[i-1]. In Tree DP, we look downward to our children. Because a parent node cannot make a decision without knowing the optimal answers of its children, we MUST use Post-Order Traversal (Left, Right, Root).

The Standard Template:

1. Base Case: If node == null, return the default state (usually 0 or a [0, 0] array).
2. Recursively call the function on the left child to get the left states.
3. Recursively call the function on the right child to get the right states.
4. Calculate the current node's state using the left and right states.
5. Return the current state to the parent.

3. Beginner Code Example: Diameter of Binary Tree

Problem Statement: (LeetCode 543) Given the root of a binary tree, return the length of the diameter of the tree. The diameter is the length of the longest path between any two nodes. This path may or may not pass through the root.

This perfectly illustrates the core issue of Tree DP: A path can either go through the current node (bridging the left and right sides) or the optimal path might be entirely contained within one of the children.

The Insight:

• A path going through the node has length: max_depth(left) + max_depth(right).
• We maintain a global variable to track the maximum diameter found across ALL nodes.
• The recursive function simply returns the max_depth to its parent so the parent can calculate its own diameter.

class TreeNode {
    int val;
    TreeNode left, right;
}

public class TreeDiameter {
    private int maxDiameter = 0;
    
    public int diameterOfBinaryTree(TreeNode root) {
        calculateDepth(root);
        return maxDiameter;
    }
    
    // Returns the max depth of the subtree rooted at 'node'
    private int calculateDepth(TreeNode node) {
        if (node == null) return 0; // Base case
        
        // 1. Get states from children
        int leftDepth = calculateDepth(node.left);
        int rightDepth = calculateDepth(node.right);
        
        // 2. Compute the local diameter passing THROUGH this node
        int localDiameter = leftDepth + rightDepth;
        
        // 3. Update the global maximum diameter
        maxDiameter = Math.max(maxDiameter, localDiameter);
        
        // 4. Return the depth to the parent (Pick the deeper side + 1 for current node)
        return Math.max(leftDepth, rightDepth) + 1;
    }
}

Complexity Analysis:

• Time Complexity: $O(N)$. We visit every node exactly once.
• Space Complexity: $O(H)$ for the recursion stack, where $H$ is the height of the tree.

4. Intermediate Code Example: Binary Tree Maximum Path Sum

Problem Statement: (LeetCode 124) Find the maximum path sum between any two nodes. The nodes can have negative values!

This is a much harder version of the Diameter problem. Because values can be negative, a path might choose to completely ignore a child subtree if its sum is negative!

The Insight:

• We only want to include a child's path if it contributes a POSITIVE sum. If a child returns a negative sum, we treat it as 0 (we sever the path).
• The maximum path passing through the current node is node.val + max(0, leftSum) + max(0, rightSum).
• The value we return to the parent can only be a single straight line! A parent cannot use a path that branches into both of our children. So we return node.val + max(leftSum, rightSum).

public class MaxPathSum {
    private int globalMax = Integer.MIN_VALUE;
    
    public int maxPathSum(TreeNode root) {
        calculateMaxPath(root);
        return globalMax;
    }
    
    private int calculateMaxPath(TreeNode node) {
        if (node == null) return 0;
        
        // Post-order traversal. Ignore negative paths!
        int leftSum = Math.max(0, calculateMaxPath(node.left));
        int rightSum = Math.max(0, calculateMaxPath(node.right));
        
        // Calculate the max path sum WITH a split at this node
        int currentPathSum = node.val + leftSum + rightSum;
        
        // Update global max
        globalMax = Math.max(globalMax, currentPathSum);
        
        // Return the maximum straight-line path to the parent
        return node.val + Math.max(leftSum, rightSum);
    }
}

5. Advanced Code Example: House Robber III

Problem Statement: (LeetCode 337) The thief has found a new neighborhood. It forms a binary tree. If two directly-linked houses are broken into on the same night, the police are alerted. Determine the maximum amount of money the thief can rob without alerting the police.

This is the ultimate Tree DP problem. We cannot just use a global variable. For every node, the parent needs to know two different pieces of information:

1. The max money if we ROB this node.
2. The max money if we SKIP this node.

We return an int[] of size 2 from our recursive function! [Rob, Skip].

The Insight:

• If we ROB the current node, we CANNOT rob its children. So we must add the children's Skip values. current[0] = node.val + left[1] + right[1].
• If we SKIP the current node, the children are free to be robbed OR skipped (whichever yields more money!). current[1] = max(left[0], left[1]) + max(right[0], right[1]).

public class HouseRobberIII {
    public int rob(TreeNode root) {
        int[] result = robSub(root);
        // Return the max of robbing the root vs skipping the root
        return Math.max(result[0], result[1]);
    }
    
    // Returns an array: [Rob Current Node, Skip Current Node]
    private int[] robSub(TreeNode node) {
        if (node == null) {
            return new int[]{0, 0}; // Base case
        }
        
        // Post-order: Get states from children
        int[] left = robSub(node.left);
        int[] right = robSub(node.right);
        
        int[] current = new int[2];
        
        // Case 1: ROB current node (Must skip children)
        current[0] = node.val + left[1] + right[1];
        
        // Case 2: SKIP current node (Can pick the max of robbing or skipping children)
        current[1] = Math.max(left[0], left[1]) + Math.max(right[0], right[1]);
        
        return current;
    }
}

Complexity Analysis:

• Time Complexity: $O(N)$. We visit every node exactly once, doing $O(1)$ math per node.
• Space Complexity: $O(H)$ for the recursion stack. We don't need a massive $O(N)$ Hash Map cache because the return values directly pass the state up the tree!

6. Common Mistakes Beginners Make

1. Trying to use a HashMap Cache (Memoization)

   • The Mistake: Writing a Top-Down DP solution with a Map<TreeNode, Integer> memo.
   • Why it's wrong: While technically correct, computing the hash of a custom object is slow, and managing the cache is clunky.
   • The Fix: Use the "Return State Array" pattern (like in House Robber III). It is vastly cleaner and achieves true $O(N)$ time with minimal overhead.

2. Returning the "Split" path to the parent

   • The Mistake: In Maximum Path Sum, returning node.val + leftSum + rightSum.
   • Why it's wrong: A path in a tree cannot fork in two directions and then continue upwards. It must be a continuous single line. A parent can only connect to ONE of your children through you. You must return node.val + max(left, right).

3. Pre-Order Traversal

   • The Mistake: Doing calculations before calling recurse(left).
   • Why it's wrong: Tree DP relies strictly on Bottom-Up information flow. You cannot make an optimal DP decision if you don't know the optimal DP states of your subtrees. Always use Post-Order.

7. Interview Q&A

Q1: What if the tree is not a Binary Tree, but an N-ary Tree (or a Graph)? Answer: The exact same principles apply! Instead of left and right, you loop over the node.children list. You aggregate the states of all children (e.g., using a for loop to sum up the values) and return the state to the parent.

Q2: (LeetCode 968) Binary Tree Cameras: How do you handle states when a node affects its parent and its children? Answer: This requires a 3-state return array!

1. Node has a camera. 2) Node is covered by a camera (but doesn't have one). 3) Node is NOT covered (requires the parent to place a camera). You analyze all permutations of the children's 3 states to determine the parent's optimal state.

Q3: Can Tree DP be implemented Iteratively? Answer: Yes, using a Post-Order iterative traversal with a Stack. However, it is extraordinarily difficult to track the return states of the children without attaching extra metadata to the Nodes or using HashMaps. In interviews, Tree DP should almost exclusively be solved using Recursion.

8. Practice Problems

1. Diameter of Binary Tree (LeetCode 543)
   • Goal: Master the global variable vs return value separation.
2. Binary Tree Maximum Path Sum (LeetCode 124)
   • Goal: Master pruning negative branches with Math.max(0, val).
3. House Robber III (LeetCode 337)
   • Goal: Master the int[] return state pattern.

9. Summary

• Post-Order is Mandatory: The left and right subtrees must be fully processed before the root can make a decision.
• Two Communication Channels:
  1. Global Variable: Tracks the absolute best answer found anywhere in the tree (often involves paths that split at a node).
  2. Return Value: The strict, straight-line optimal state that is passed back to the parent to continue the recursion.
• State Arrays: If a node has multiple choices (e.g., Rob vs Skip), return an array representing the optimal cost for each choice.