Dynamic programming on trees

You’ll recognise it when

• The input is a tree: a TreeNode, a parent array, or n - 1 edges. You want one best value or count over it, or one per node.
• A node’s answer depends only on its own value and its children’s answers: subtree sizes, heights, sums, “is this subtree balanced?”.
• The question is about paths that can start and end anywhere: the longest path, or the path with the largest sum.
• You pick or label nodes under a rule about neighbours: never a parent and its child together, every node watched by a neighbour, neighbours in different colours.
• Trying every start node, or every subset of nodes, is far too slow.

It looks like a plain post-order traversal, and it is one. The DP part is deciding what each call returns.

The idea

Think of a company that wants one number from the whole organisation, like the largest team. Nobody reads the full org chart. Each manager asks each direct report for a short summary, adds their own part, and passes one summary up to their own boss. Every person is asked once, and the answer arrives at the top.

That’s tree DP. A subtree is a smaller problem of the same kind, and different subtrees never overlap. So you solve the children first and build each node’s answer from theirs. The DP table has one entry per node, and post-order is the order that fills it.

The hard part is choosing what to report. Ask: what does my parent need from me? Often it isn’t the final answer. Take the largest sum of any path, where a path may start and end at any nodes and values can be negative:

• Every path has one highest node, where it bends. There, the path is the node plus the best path going straight down into its left side, plus the same on its right. Either side can be left out if it would lower the sum.
• The parent can only use a path that goes straight down from the node. A path that has already bent at the node can’t also go up to the parent: the node would have three neighbours on it.

So each call returns down, the best path that starts at the node and goes straight down, and on the side it checks bend against a global best. The tree traversal article shows this shape for the diameter. Here it is for path sums:

def max_path_sum(root):
    best = root.val
    def down(node):            # best path straight down
        nonlocal best
        if node is None:
            return 0
        left = max(down(node.left), 0)   # drop a losing side
        right = max(down(node.right), 0)
        # the path that bends here
        best = max(best, node.val + left + right)
        # what the parent can extend
        return node.val + max(left, right)
    down(root)
    return best

How it works

The code in the tabs does the same work without recursion, on a tree with some negative values, so some sides are worth dropping. Scroll through the steps and the graphic follows along. You can also turn on quiz me to predict what each node hands up, or edit the tree.

Why it’s right: every path in a tree has exactly one highest node. At that node, the path is the node plus at most one downward path into each child, so its best possible sum is that node’s bend, as long as each child’s down is right. down is right by induction from the leaves: a downward path is either the node alone, or the node plus a downward path from one child. Every node’s bend is checked against best, so the best path can’t be missed.

class TreeNode:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def max_path_sum(root):
    """Largest sum of any non-empty path, in one bottom-up pass."""
    # Parents before children; read backwards, children come first.
    order, stack = [], [root]
    while stack:
        node = stack.pop()
        order.append(node)
        for child in (node.left, node.right):
            if child:
                stack.append(child)

    down = {None: 0}    # best path from a node straight down; an empty child gives 0
    best = root.val     # one node on its own is a path
    for node in reversed(order):
        left = max(down[node.left], 0)            # @step:children  a losing branch counts as 0
        right = max(down[node.right], 0)
        bend = node.val + left + right            # @step:bend      the best path turning here
        best = max(best, bend)
        down[node] = node.val + max(left, right)  # @step:return    the parent extends one side only
    return best

#include <algorithm>
#include <unordered_map>
#include <vector>
using namespace std;

struct TreeNode {
    int val;
    TreeNode* left = nullptr;
    TreeNode* right = nullptr;
    explicit TreeNode(int v) : val(v) {}
};

// Largest sum of any non-empty path, in one bottom-up pass.
long long maxPathSum(TreeNode* root) {
    // Parents before children; read backwards, children come first.
    vector<TreeNode*> order, stack = {root};
    while (!stack.empty()) {
        TreeNode* node = stack.back();
        stack.pop_back();
        order.push_back(node);
        for (TreeNode* child : {node->left, node->right})
            if (child) stack.push_back(child);
    }

    // Best path from a node straight down; an empty child gives 0.
    unordered_map<TreeNode*, long long> down = {{nullptr, 0}};
    long long best = root->val;                           // one node on its own is a path
    for (auto it = order.rbegin(); it != order.rend(); ++it) {
        TreeNode* node = *it;
        long long left = max(down[node->left], 0LL);      // @step:children  a losing branch counts as 0
        long long right = max(down[node->right], 0LL);
        long long bend = node->val + left + right;        // @step:bend      the best path turning here
        best = max(best, bend);
        down[node] = node->val + max(left, right);        // @step:return    the parent extends one side only
    }
    return best;
}

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int val) { this.val = val; }
}

class MaxPathSum {
    // Largest sum of any non-empty path, in one bottom-up pass.
    static long maxPathSum(TreeNode root) {
        // Parents before children; read backwards, children come first.
        List<TreeNode> order = new ArrayList<>();
        Deque<TreeNode> stack = new ArrayDeque<>();
        stack.push(root);
        while (!stack.isEmpty()) {
            TreeNode node = stack.pop();
            order.add(node);
            if (node.left != null) stack.push(node.left);
            if (node.right != null) stack.push(node.right);
        }

        // Best path from a node straight down; an empty child gives 0.
        Map<TreeNode, Long> down = new HashMap<>();
        down.put(null, 0L);
        long best = root.val;                                // one node on its own is a path
        for (int i = order.size() - 1; i >= 0; i--) {
            TreeNode node = order.get(i);
            long left = Math.max(down.get(node.left), 0);    // @step:children  a losing branch counts as 0
            long right = Math.max(down.get(node.right), 0);
            long bend = node.val + left + right;             // @step:bend      the best path turning here
            best = Math.max(best, bend);
            down.put(node, node.val + Math.max(left, right)); // @step:return   the parent extends one side only
        }
        return best;
    }
}

The dictionary down holds what the recursive calls would have returned. It is the DP table: one entry per node, each filled once, after its children. The recursive version is shorter, but a tree shaped like a long chain makes it one call deeper per node, and Python stops at about 1,000 calls. sys.setrecursionlimit lets it go deeper, but a very deep recursion can then crash the interpreter itself. C++ and Java go much deeper, but a chain of a million nodes can overflow their stacks too. The explicit order has no depth limit.

Sometimes one number isn’t enough. In DP basics, you rob houses on a street, never two neighbours. Put the houses on a tree, where a parent and its child are neighbours. Whether a node can be robbed now depends on its parent’s choice, which the node can’t see. So each call returns two numbers, one for each choice, and lets the parent decide:

def rob(node):
    """(best if node is robbed, best if node is skipped)"""
    if node is None:
        return 0, 0
    l_rob, l_skip = rob(node.left)
    r_rob, r_skip = rob(node.right)
    take = node.val + l_skip + r_skip    # children skipped
    skip = max(l_rob, l_skip) + max(r_rob, r_skip)  # free
    return take, skip

The answer is max(rob(root)). The same trick covers any rule about neighbours: return one number per state the parent cares about, like one count per colour, or “has a camera, watched, not watched”.

Why it’s O(n)

Each node is visited once. In a binary tree it does a constant amount of work with its children’s summaries. In a general tree it does work in proportion to its children, and all the children add up to n - 1. So the whole pass is O(n).

A node’s answer covers its whole subtree, but it only reads its children’s numbers, never the subtree itself. That’s where the saving comes from. The obvious alternative, starting a search from every node, repeats almost everything:

Common mistakes

return node.val + left + right      # ✗ a fork
return node.val + max(left, right)  # ✓ one side

left = down[node.left]          # ✗ drags it down
left = max(down[node.left], 0)  # ✓ drop it

best = 0          # ✗ 0 for [-3, -1, -8]
best = root.val   # ✓ -1 for [-3, -1, -8]

abs(height(n.left) - height(n.right)) <= 1   # ✗ walks again
if l < 0 or r < 0 or abs(l - r) > 1:         # ✓ one pass
    return -1
return 1 + max(l, r)

Variations

• Diameter. The same shape with lengths instead of sums: return the height, and check the two tallest child heights added together. In a tree with any number of children, keep the two largest values from the children.
• Balanced check. Return the height, or -1 as soon as any subtree is out of balance. A node is balanced if both children are and their heights differ by at most 1.
• Trees given as edges or parents. Root the tree at any node. Get a parents-first order with a stack or BFS, recording each node’s parent, then walk it backwards and add each node’s summary into its parent’s: size[parent[v]] += size[v] counts everyone below each node.
• Several states per node. Return one number per state: robbed or skipped, one count per colour, or “has a camera, watched, not watched”. The parent combines the children’s states under its rule.
• Every node as the root. When the question asks for an answer with each node as the root, solve it for one root, then run a second pass from the top that moves the root across each edge and fixes up the answer in O(1). That’s rerooting: two passes, O(n), instead of one pass per root.