🎓 A1Learner
🔍
0 of 93 lessons completed

DSA The Tree DP Pattern

Tree dynamic programming solves a tree problem by returning a small tuple of numbers from every node to its parent. Each tuple summarises a whole subtree in exactly the facts the parent needs.

First design: count and height

The empty child returns (0, -1), which adds no node and no height. The root of this small tree reports three nodes and a height of one edge.

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

def measure(node):
    if node is None:
        return (0, -1)
    left = measure(node.left)
    right = measure(node.right)
    count = 1 + left[0] + right[0]
    height = 1 + max(left[1], right[1])
    return (count, height)


tree = TreeNode(8, TreeNode(3), TreeNode(10))
print("Pair (count, height):", measure(tree))

Second design: total and smallest

The same tree is measured by a different tuple here. The empty child returns (0, None), and None is ignored when taking the smallest value.

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

def measure(node):
    if node is None:
        return (0, None)
    left = measure(node.left)
    right = measure(node.right)
    total = node.value + left[0] + right[0]
    smallest = node.value
    for part in (left[1], right[1]):
        if part is not None:
            smallest = min(smallest, part)
    return (total, smallest)


tree = TreeNode(8, TreeNode(3), TreeNode(10))
print("Pair (total, smallest):", measure(tree))

Example

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

def count_height(node):
    if node is None:
        return (0, -1)
    left = count_height(node.left)
    right = count_height(node.right)
    return (1 + left[0] + right[0], 1 + max(left[1], right[1]))


def total_smallest(node):
    if node is None:
        return (0, None)
    left = total_smallest(node.left)
    right = total_smallest(node.right)
    total = node.value + left[0] + right[0]
    smallest = node.value
    for part in (left[1], right[1]):
        if part is not None:
            smallest = min(smallest, part)
    return (total, smallest)


def report(node, rows):
    if node is None:
        return
    report(node.left, rows)
    report(node.right, rows)
    rows.append((node.value, count_height(node), total_smallest(node)))


tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6)),
                TreeNode(10, None, TreeNode(14)))

rows = []
report(tree, rows)
for value, first, second in rows:
    print("Node", value, "count-height:", first, "total-smallest:", second)
print("Root, first design:", count_height(tree))
print("Root, second design:", total_smallest(tree))

Key points

  • The design step is choosing the tuple, the few numbers that describe a subtree's answer
  • Children return their tuples first, and the parent combines them with its own value
  • An empty child returns an identity tuple that changes nothing in the combination
  • The final answer is the root tuple itself, or a best value tracked during the walk
  • One tree can carry two different tuple designs, each answering a different question

In this example, one six-node tree is measured by both tuple designs. The rows print in postorder, so every child row appears before its parent row. The root row shows a count of 6 and a height of 2 under the first design. Under the second design the same root shows a total of 42 and a smallest value of 1.

Note

The maximum path sum and House Robber III lessons in this entry are both this pattern with different tuples

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In a tree DP design, what must the tuple returned by an empty child do?