Dynamic Programming on Trees

Dynamic programming works whenever a problem decomposes into overlapping subproblems ordered so that each can be solved from smaller ones already in hand. On a sequence the natural subproblems are prefixes; on an interval they are subintervals. On a tree the natural subproblems are rooted subtrees, and the ordering that makes them solvable is the post-order traversal, which visits every node only after all of its children. Root the tree anywhere, define an answer $f(v)$ for the subtree hanging below each node $v$, and a single depth-first sweep fills the entire table.

The defining feature of these problems is that the recurrence is local: the answer at $v$ depends only on the answers at $v$'s children,

never on grandchildren directly and never on the rest of the tree. Because the DFS touches each node and each edge exactly once and does $O(1)$ work per child, the whole computation runs in $\Theta (n)$ time for a tree on $n$ nodes.¹ The work is in choosing the state: what $f(v)$ must remember about the subtree so that a parent can combine children without re-descending into them. This is the usual optimal-substructure question, specialized to subtrees.

The archetype: maximum-weight independent set

Let each node $v$ of a tree carry a weight $w_v\ge 0$. An independent set is a set of nodes no two of which are adjacent; we want one of maximum total weight. On a general graph this is NP-hard, but on a tree dynamic programming solves it in linear time, the canonical illustration of the whole technique.²

The idea is to make the state record whether $v$ itself is used, because that is the one fact a parent needs in order to decide about itself. Define, for the subtree rooted at $v$, two values:

• $dp[v][0]$, the best independent set of $v$'s subtree in which $v$ is not taken;
• $dp[v][1]$, the best one in which $v$ is taken.

If $v$ is not taken, each child $c$ is free to be taken or not, so we keep the better of its two options. If $v$ is taken, no child may be taken, so each child must contribute its $dp[c][0]$:

The base case falls out for free: a leaf has no children, so $dp[v][0]=0$ and $dp[v][1]=w_v$. The answer for the whole tree is $\max \left(dp[\text{root}][0],dp[\text{root}][1]\right)$. This is House Robber III, where weights are the money in each house and the adjacency constraint forbids robbing a parent and its child on the same night.

The procedure visits each node once and spends $O(1)$ per child, hence $\Theta (n)$ total. The state, taken vs. not taken, is the part worth remembering: a single extra bit per node turns an intractable graph problem into a linear-time tree sweep.³

To see the post-order fill in full, take the tree in the figure with root $r$ (weight $5$), children $b$ (weight $3$) and $c$ (weight $4$), then leaves $d,e$ (weights $3,1$) under $b$ and leaf $f$ (weight $1$) under $c$. Post-order visits the leaves first, then $b$, then $c$, then $r$. Each row is the pair $(dp[v][0],dp[v][1])=(\text{skip }v,\text{take }v)$:

The answer is $\max (dp[r][0],dp[r][1])=\max (8,10)=10$, achieved by taking $r$. Taking $r$ forbids $b$ and $c$, so we descend into $dp[b][0]$ and $dp[c][0]$, each of which skips its own node and is free to take the leaves below: that selects $r$, $d$, $e$, $f$, with weights $5+3+1+1=10$ — the shaded set in the figure. Every value in the table is read from children already computed, never recomputed — that single-pass reuse is what makes the sweep linear.

from __future__ import annotations

from typing import Generic, NamedTuple, Optional, TypeVar

Weight = TypeVar("Weight", int, float)

class TreeNode(Generic[Weight]):
  """
    A rooted-tree node: its own weight and its list of child nodes.\n
  """

  def __init__(
    self,
    weight: Weight,
    children: Optional[list[TreeNode[Weight]]] = None,
  ) -> None:
    self.weight: Weight = weight
    self.children: list[TreeNode[Weight]] = children if children is not None else []

  def __repr__(self) -> str:
    return f"TreeNode(weight={self.weight!r}, children={len(self.children)})"

class SubtreeChoice(NamedTuple):
  """
    The two best subtree weights: with the node skipped, and with it taken.\n
  """
  skip: float
  take: float

def _solve(node: TreeNode[Weight]) -> SubtreeChoice:
  """
    Post-order combine returning the (skip, take) pair for `node`'s subtree.\n
    Skipping the node lets each child take its better option; taking the\n
    node forbids every child, so each must contribute its own skip value.\n
  """
  # skipping frees children to choose; taking adds this weight, bars children.
  skip_total: float = 0.0
  take_total: float = float(node.weight)

  # each child contributes its better option when skipped, its skip when taken.
  for child in node.children:
    child_choice: SubtreeChoice = _solve(child)
    skip_total += max(child_choice.skip, child_choice.take)
    take_total += child_choice.skip

  return SubtreeChoice(skip_total, take_total)

def max_weight_independent_set(root: Optional[TreeNode[Weight]]) -> float:
  """
    Maximum total weight of an independent set in the tree at `root`.\n
    Returns 0 for an empty tree.\n
  """
  if root is None:
    return 0.0
  choice: SubtreeChoice = _solve(root)
  return max(choice.skip, choice.take)

Paths through a node: diameter and maximum path sum

A second pattern arises when the quantity we care about is a path, not a set. The diameter of a tree is the number of edges on its longest path; the maximum path sum (where nodes carry values, possibly negative) is the largest total along any path. Neither is a clean subtree quantity, because the optimal path may bend at some node, descending into two different children.

The resolution is the signature move of tree DP on paths. At each node $v$, let $\text{down}(v)$ be the best downward path that starts at $v$ and goes into a single subtree. A child $c$ extends to $\text{down}(c)+(\text{edge or }{w}_v)$. The best path that bends at $v$ combines its two best children:

for the two children with the largest downward values. The distinction below is the source of the classic bug:

For Binary Tree Maximum Path Sum the $\max (0,\cdot )$ prunes branches that would only hurt the total; the global $ans$ records the best bend seen anywhere. On the tree in the figure — root $-10$ with left leaf $9$ and right child $20$ whose children are leaves $15$ and $7$ — the post-order pass runs as follows. The leaves $9$, $15$, $7$ each return their own value (no children), and update $ans$ with themselves. At node $20$: $L=\max (0,15)=15$, $R=\max (0,7)=7$, so the bent path is $20+15+7=42$, which becomes the new $ans$; it returns upward only $20+\max (15,7)=35$. At the root $-10$: $L=\max (0,9)=9$, $R=\max (0,35)=35$, and its bent path is $-10+9+35=34<42$, so $ans$ stays $42$. The negative root could not improve the answer, and the $\max (0,\cdot )$ guards ensured no negative branch was added into a sum. The final answer is $42$, the two-child bend at $20$ — a path that the node correctly updated the global with but did not return.

For Diameter of Binary Tree the same skeleton applies with edge counts in place of values: $\text{down}(v)=1+\max (\text{down}(\text{children}))$ and the diameter is the largest $\text{down}(c_1)+\text{down}(c_2)$ over all nodes $v$. Both run in $\Theta (n)$: one post-order pass, $O(1)$ per node.

from __future__ import annotations

from typing import Optional

class TreeNode:
  """
    A binary-tree node: a value and optional left and right children.\n
  """

  def __init__(
    self,
    value: int,
    left: Optional[TreeNode] = None,
    right: Optional[TreeNode] = None,
  ) -> None:
    self.value: int = value
    self.left: Optional[TreeNode] = left
    self.right: Optional[TreeNode] = right

  def __repr__(self) -> str:
    return f"TreeNode(value={self.value!r})"

def max_path_sum(root: Optional[TreeNode]) -> int:
  """
    The maximum sum over any non-empty path in the tree at `root`.\n
    Raises ValueError on an empty tree (a path must contain a node).\n
  """
  if root is None:
    raise ValueError("max_path_sum requires a non-empty tree")
  best: int = root.value

  def best_downward(node: Optional[TreeNode]) -> int:
    """
      The largest sum of a downward path starting at `node`, after dropping\n
      any child branch whose contribution is negative; meanwhile updates the\n
      global best with the path that bends through `node` using both sides.\n
    """
    nonlocal best
    if node is None:
      return 0

    # recurse into each side, pruning any branch that would only subtract.
    left_gain: int = max(0, best_downward(node.left))
    right_gain: int = max(0, best_downward(node.right))

    # the bend here joins both branches; it can never extend upward.
    best = max(best, node.value + left_gain + right_gain)

    # only one branch may extend into the parent.
    return node.value + max(left_gain, right_gain)

  best_downward(root)
  return best

from __future__ import annotations

from typing import Optional

class TreeNode:
  """
    A rooted-tree node holding only its child links (the diameter counts\n
    edges, so node values are irrelevant).\n
  """

  def __init__(self, children: Optional[list[TreeNode]] = None) -> None:
    self.children: list[TreeNode] = children if children is not None else []

  def __repr__(self) -> str:
    return f"TreeNode(children={len(self.children)})"

def tree_diameter(root: Optional[TreeNode]) -> int:
  """
    The diameter (longest path length in edges) of the tree at `root`.\n
    An empty tree and a single node both have diameter 0.\n
  """
  longest_path: int = 0

  def depth(node: TreeNode) -> int:
    """
      Edges on the deepest downward path from `node`, updating the diameter\n
      with the best bend (two longest child depths) through this node.\n
    """
    nonlocal longest_path
    best_down: int = 0
    second_down: int = 0

    # track the two deepest downward branches among the children.
    for child in node.children:
      child_down: int = depth(child) + 1
      if child_down > best_down:
        best_down, second_down = child_down, best_down
      else:
        second_down = max(second_down, child_down)

    # the path bending here joins the two deepest branches.
    longest_path = max(longest_path, best_down + second_down)
    return best_down

  # empty tree has no path; otherwise the DFS fills longest_path.
  if root is None:
    return 0
  depth(root)
  return longest_path

Rerooting: an answer for every root in $O(n)$

The hardest variant asks for a quantity computed with each node in turn as the root: for every node $v$, say, the sum of distances from $v$ to all other nodes. Re-running an $O(n)$ DFS from each of the $n$ roots costs $O(n^2)$. Rerooting (also called the all-roots or re-root technique) computes all $n$ answers in $O(n)$ total, with two DFS passes: one down, one up.⁴

Take Sum of Distances in Tree. Fix an arbitrary root $r$ and let $S(v)$ be the number of nodes in $v$'s subtree and $D(v)$ the sum of distances from $v$ to every node inside its own subtree. A post-order pass computes both, since a child $c$ at distance $1$ contributes $D(c)+S(c)$ (every node under $c$ is one edge farther from $v$ than from $c$):

That gives the true global answer only at the root, where the subtree is the whole tree: $\text{ans}(r)=D(r)$. The second pass pushes the answer from a parent to each child in $O(1)$. Moving the root from $u$ to an adjacent child $v$, the $S(v)$ nodes on $v$'s side each get one closer (distance drops by $1$) and the remaining $n-S(v)$ nodes each get one farther:

Subtract the subtree's contribution, add the rest: the entire adjustment is a single $O(1)$ formula, so the down-pass plus the up-pass together are $\Theta (n)$.

On the five-node tree above (root $r$ with children $a,b$; then $c$ under $a$ and $d$ under $b$), the down-pass fixes $\text{ans}(r)=D(r)=6$. The up-pass then propagates outward, each step a single subtract-add. Moving to $a$: its subtree has $S(a)=2$ nodes, so $\text{ans}(a)=6-2+(5-2)=7$. From $a$ to its child $c$ ($S(c)=1$): $\text{ans}(c)=7-1+(5-1)=10$. By symmetry $\text{ans}(b)=7$ and $\text{ans}(d)=10$. Every value matches a direct BFS from that node, but the whole sweep is linear.

A related linear-time tree DP is Distribute Coins in Binary Tree: each node returns to its parent the net coins it must send up or pull down, the signed excess $\text{coins}-1$ summed over its subtree, and the total number of moves is the sum of absolute flows along every edge, accumulated in one post-order pass. Same shape: a local return value, a global accumulator, $\Theta (n)$ time.

from collections.abc import Hashable
from typing import TypeVar

from graph import Graph, Vertex

Label = TypeVar("Label", bound=Hashable)

def sum_of_distances(tree: Graph[Label]) -> dict[Label, int]:
  """
    Map each node label to the sum of distances (edge counts) from that node\n
    to all others, computing every node's answer in O(n) by rerooting.\n
    `tree` must be an undirected, connected, acyclic graph. A single node\n
    maps to 0; an empty graph maps to an empty dict.\n
  """
  # an empty graph has no answers to compute.
  node_count: int = len(tree)
  if node_count == 0:
    return {}

  # caches filled by the down-pass; answers filled by the up-pass.
  subtree_size: dict[Label, int] = {}
  distance_sum: dict[Label, int] = {}
  answer: dict[Label, int] = {}

  # root the tree arbitrarily at the first vertex.
  root: Vertex[Label] = tree.vertices[0]

  # down-pass: discover a parent-pointed order with an explicit stack.
  visited_down: set[Label] = {root.label}
  ordering: list[Vertex[Label]] = []
  parent_of: dict[Label, Label] = {}
  stack: list[Vertex[Label]] = [root]

  # flood outward, recording discovery order and each node's parent.
  while stack:
    current: Vertex[Label] = stack.pop()
    ordering.append(current)

    # push each undiscovered neighbor, remembering current as its parent.
    for neighbor in current.neighbors():
      if neighbor.label not in visited_down:
        visited_down.add(neighbor.label)
        parent_of[neighbor.label] = current.label
        stack.append(neighbor)

  # process children before parents by walking discovery order in reverse.
  for vertex in reversed(ordering):
    size_here: int = 1
    distance_here: int = 0

    # fold in each child: its size, plus one extra edge per node under it.
    for neighbor in vertex.neighbors():
      if parent_of.get(neighbor.label) == vertex.label:
        size_here += subtree_size[neighbor.label]
        distance_here += distance_sum[neighbor.label] + subtree_size[neighbor.label]

    subtree_size[vertex.label] = size_here
    distance_sum[vertex.label] = distance_here

  # the root's in-subtree sum spans the whole tree, so it is the true answer.
  answer[root.label] = distance_sum[root.label]

  # up-pass: pre-order from the root, deriving each child's answer.
  visited_up: set[Label] = {root.label}
  frontier: list[Vertex[Label]] = [root]
  while frontier:
    current = frontier.pop()

    # moving the root to a child: its side gets 1 closer, the rest 1 farther.
    for child in current.neighbors():
      if child.label in visited_up:
        continue
      visited_up.add(child.label)
      closer: int = subtree_size[child.label]

      answer[child.label] = answer[current.label] - closer + (node_count - closer)
      frontier.append(child)

  return answer

from __future__ import annotations

from typing import Optional

class TreeNode:
  """
    A binary-tree node: the coins it holds and optional left/right children.\n
  """

  def __init__(
    self,
    coins: int,
    left: Optional[TreeNode] = None,
    right: Optional[TreeNode] = None,
  ) -> None:
    self.coins: int = coins
    self.left: Optional[TreeNode] = left
    self.right: Optional[TreeNode] = right

  def __repr__(self) -> str:
    return f"TreeNode(coins={self.coins!r})"

def distribute_coins(root: Optional[TreeNode]) -> int:
  """
    The minimum number of single-coin moves that leaves exactly one coin at\n
    every node. Assumes total coins equals the node count. Empty tree: 0.\n
  """
  total_moves: int = 0

  def excess(node: Optional[TreeNode]) -> int:
    """
      The net coins `node`'s subtree must exchange with its parent: positive\n
      means it ships coins up, negative means it pulls coins down. Each child\n
      edge carries abs(child excess) moves regardless of direction.\n
    """
    nonlocal total_moves
    if node is None:
      return 0

    # each edge to a child carries abs(excess) moves, either direction.
    left_excess: int = excess(node.left)
    right_excess: int = excess(node.right)
    total_moves += abs(left_excess) + abs(right_excess)

    # this subtree's surplus (or deficit): its coins minus the one it keeps.
    return node.coins - 1 + left_excess + right_excess

  excess(root)
  return total_moves

from collections.abc import Hashable, Iterator
from typing import Generic, Optional, TypeVar


Label = TypeVar("Label", bound=Hashable)


class Edge(Generic[Label]):
  """
    A directed connection from `source` to `target`, carrying a weight.\n
  """

  def __init__(
    self,
    source: Vertex[Label],
    target: Vertex[Label],
    weight: float = 1.0,
  ) -> None:
    self.source: Vertex[Label] = source
    self.target: Vertex[Label] = target
    self.weight: float = weight

  def __repr__(self) -> str:
    return f"Edge({self.source.label!r} -> {self.target.label!r}, w={self.weight})"


class Vertex(Generic[Label]):
  """
    A graph vertex: a label plus the list of edges leaving it.\n
  """

  def __init__(self, label: Label) -> None:
    self.label: Label = label
    self.outgoing: list[Edge[Label]] = []

  def neighbors(self) -> list[Vertex[Label]]:
    """
      The vertices reachable from this one by a single edge.\n
    """
    return [edge.target for edge in self.outgoing]

  def edge_to(self, label: Label) -> Optional[Edge[Label]]:
    """
      The outgoing edge to the vertex with `label`, or None.\n
    """
    for edge in self.outgoing:
      if edge.target.label == label:
        return edge
    return None

  def __repr__(self) -> str:
    return f"Vertex({self.label!r})"


class Graph(Generic[Label]):
  """
    A graph of Vertex objects linked by Edge objects.\n
    Pass `directed=True` for a digraph; otherwise each `add_edge` inserts\n
    the reverse edge too.\n
  """

  def __init__(self, directed: bool = False) -> None:
    self.directed: bool = directed
    self._vertices: dict[Label, Vertex[Label]] = {}

  def add_vertex(self, label: Label) -> Vertex[Label]:
    """
      Return the vertex for `label`, creating it if it is absent.\n
    """
    # reuse the existing vertex, or mint and register a fresh one.
    vertex = self._vertices.get(label)
    if vertex is None:
      vertex = Vertex(label)
      self._vertices[label] = vertex
    return vertex

  def add_edge(
    self,
    source_label: Label,
    target_label: Label,
    weight: float = 1.0,
  ) -> None:
    """
      Connect two labels (creating either vertex as needed).\n
      Adds the reverse edge as well when the graph is undirected.\n
    """
    source = self.add_vertex(source_label)
    target = self.add_vertex(target_label)

    # link source to target, and mirror it back when undirected.
    source.outgoing.append(Edge(source, target, weight))
    if not self.directed:
      target.outgoing.append(Edge(target, source, weight))

  def vertex(self, label: Label) -> Vertex[Label]:
    """
      The vertex carrying `label` (raises KeyError if absent).\n
    """
    return self._vertices[label]

  @property
  def vertices(self) -> list[Vertex[Label]]:
    """
      Every vertex, in insertion order.\n
    """
    return list(self._vertices.values())

  @property
  def labels(self) -> list[Label]:
    """
      Every vertex label, in insertion order.\n
    """
    return list(self._vertices)

  def edges(self) -> Iterator[Edge[Label]]:
    """
      Each edge once — an undirected edge is yielded a single time.\n
    """
    # track undirected endpoint pairs so each is emitted only once.
    seen: set[frozenset[Label]] = set()

    for vertex in self._vertices.values():
      for edge in vertex.outgoing:
        # skip an undirected edge already yielded from the other endpoint.
        if not self.directed:
          endpoints = frozenset((edge.source.label, edge.target.label))
          if endpoints in seen:
            continue
          seen.add(endpoints)

        yield edge

  def __contains__(self, label: Label) -> bool:
    return label in self._vertices

  def __iter__(self) -> Iterator[Vertex[Label]]:
    return iter(self._vertices.values())

  def __len__(self) -> int:
    return len(self._vertices)

From trees back to hard graphs

Tree DP is a special case of a deeper result. The reason maximum-weight independent set is linear on trees but NP-hard on general graphs is treewidth: a tree has treewidth $1$, and Courcelle's theorem (Courcelle, 1990) says that any graph property expressible in monadic second-order logic — independent set, dominating set, Hamiltonicity, $k$-coloring for fixed $k$ — is decidable in linear time on graphs of bounded treewidth, by a dynamic program over a tree decomposition. The post-order combine in this lesson is the treewidth-$1$ instance of that DP; on a width-$w$ decomposition each bag of $\le w+1$ vertices plays the role a single node plays here, and the state grows to roughly $2^w$ per bag, so the runtime is $O(2^w\cdot n)$. This is why bounded treewidth — the graph is nearly a tree — is such a useful property of an instance.⁵

Rerooting, the two-pass all-roots technique, is the tree analog of an idea that recurs across algorithms: compute one anchored answer, then transfer it along edges with a cheap difference. The same accounting drives the all-pairs flavor of many tree problems and appears in Skiena's treatment of tree DP and in competitive references under names like in-and-out DP or up-and-down DP. It also connects to centroid decomposition: both exploit that a tree, unlike a general graph, has a balanced recursive structure that turns an apparent $O(n^2)$ over all pairs of nodes into $O(n)$ or $O(n\log n)$.

A modern practical descendant is belief propagation (Pearl, 1988) on graphical models: on a tree-structured probabilistic model, the sum-product message-passing algorithm computes exact marginals in one up-pass and one down-pass — structurally identical to rerooting, with $\max$/$\sum$ over children replaced by products of messages. The independent-set recurrence here is the hard-core model special case, and the reason inference is exact on trees but only approximate (loopy BP) on general graphs is, again, that trees have no cycles to double-count.

Takeaways

• On a tree, the natural DP subproblems are rooted subtrees, solved by a single post-order DFS that combines each node's children in $O(1)$, hence $\Theta (n)$ overall, since every node and edge is processed once.
• The state must capture just what a parent needs. For maximum-weight independent set (House Robber III) that is one bit, taken vs. not taken: $dp[v][0]={\sum }_c\max (dp[c][0],dp[c][1])$ and $dp[v][1]=w_v+{\sum }_{c}dp[c][0]$.
• The path-through-a-node pattern (diameter, max path sum) returns one thing and updates another: return the single best downward extension to the parent, but update a global max with the two-child bent path, never returning the bent path.
• Rerooting computes a per-root answer for all $n$ nodes in $O(n)$ via two passes: a down-pass fixes the root's answer from subtree aggregates, an up-pass transfers it edge-by-edge with a subtract my subtree, add the rest $O(1)$ adjustment.
• The recurring design questions are always the same: what does a node return to its parent, and what aggregate must the subtree cache so the combine stays $O(1)$.