Embercode
✕
Binary TreesLesson 17 of 19

45 min read

Tree DP: return two values per node in Python

Why it matters

Suppose a company chart is a binary tree, with a boss at the top and two people under each manager. Each person has a score. You must pick a group of people for a task, and the scores add up. The rule is that you may not pick a person and their direct manager in the same group, because the manager would then be in charge of someone they already work with. Which group has the biggest total score?

The problem is called the maximum weight independent set. A weight is the number attached to a node, here the score. An independent set is a set of nodes where no two chosen nodes are joined by a link. In a tree, the links are the lines between a parent and its children.

The simple way is to try every group. A tree with 20 nodes has about one million possible groups, and a tree with 40 nodes has about one trillion. That is too slow. Dynamic programming, usually shortened to DP, solves the problem by breaking it into smaller pieces and reusing the answer for each piece. In a tree the natural pieces are the subtrees. A subtree is a node together with everything below it. This lesson uses DP on a tree, and each node sends two numbers up to its parent instead of one.

This lesson does not cover the general idea of recursion. It assumes you can write a function that calls itself on the left and right children of a node. The first lesson of this module covers traversals, which use the same pattern.

Words you will need

  • Weight: a number attached to a node. In this lesson the weights can be positive or negative.
  • Edge: the link between a parent and a child.
  • Independent set: a set of nodes where no two nodes share an edge. In a tree, a parent and its child cannot both be in the set.
  • Subtree: a node and every node below it.
  • Dynamic programming (DP): a way to solve a problem by solving smaller versions of it, passing back each answer, and combining them.
  • max(a, b): a Python function that returns the bigger of two numbers.
  • NamedTuple: a class from the typing module that holds fixed fields. class Best(NamedTuple) with two fields is a small class with names for its two values.
  • Base case: the simplest case, which the function answers without calling itself. Here, the empty subtree.

A small example first

Take a path of three nodes. Node A has weight 3. Its child B has weight 4. B's child C has weight 5. A path is a tree in which every node has at most one child.

A (3)
  \
   B (4)
     \
      C (5)

In a path, the links are A to B and B to C. The groups that obey the rule are:

  • {A} with total 3
  • {B} with total 4
  • {C} with total 5
  • {A, C} with total 8, because A and C are not joined

The biggest total is 8, from {A, C}. The job of the DP is to find this without listing the groups.

The idea: two answers for each node

Pick one node, call it V. Its subtree has two possible states, and the parent of V needs to know both.

  • take(V) is the best total in V's subtree when V itself is in the set. If V is in the set, its children cannot be in the set, because each child is joined to V. So each child can only contribute what it gets when it is not chosen. That is the child's skip value. The formula is take(V) = weight(V) + skip(left) + skip(right).
  • skip(V) is the best total in V's subtree when V is not in the set. Then each child is free. The child can be chosen or not, whichever gives more. For each child, that is the bigger of its take and skip values. The formula is skip(V) = max(take(left), skip(left)) + max(take(right), skip(right)).

Two details matter here.

First, a missing child is a subtree with no nodes. Its take value is 0 and its skip value is 0. This is the base case. It means the formulas work the same way for a node with one child, two children, or none.

Second, the answer for the whole tree is max(take(root), skip(root)). The root is either chosen or not, and the answer is the larger total.

Why two numbers are enough

Each node's choice only affects its own children. A child's choice only affects the child's own children, and so on down. The subtrees of two siblings do not touch each other. So once you know the best total for each child under each of its two states, the best total for the parent is just the sum of those pieces.

A single number per node is not enough. The parent needs to know whether the child is chosen, because that decides whether the parent can be chosen. The best total of the child alone does not say that. The mistakes section shows a case where one number gives a wrong answer.

Worked example 1: the path from the start

Work up the path from the bottom. The root is A.

NodeWeightChildren values usedtakeskip
C5none, both 050
B4C's skip = 04 + 0 = 4max(5, 0) = 5
A3B's skip = 53 + 5 = 8max(4, 5) = 5

For C, take is its weight, 5, because C has no children. Its skip is 0, because it is left out.

For B, take is 4 plus C's skip. C's skip is 0 because if B is chosen, C cannot be. So take is 4. B's skip is the bigger of C's two values, which is 5, because C can be chosen when B is not.

For A, take is 3 plus B's skip. B's skip is 5, so A's take is 8. A's skip is the bigger of B's two values, which is max(4, 5) = 5.

The answer is max(8, 5) = 8. This matches the group {A, C} from the start.

The program in the code section below prints the same result. The numbers come from the solve function, which returns both values as a Best tuple.

The code

The whole method is short. Each call computes the two results for its own node from the results of its two children.

Python · editable

Output:

example 1 solve(a) = Best(take=8, skip=5)
example 1 answer   = 8
example 2 solve(root) = Best(take=29, skip=19)
example 2 answer      = 29
example 3 solve: Best(take=14, skip=11)
example 3 answer: 14
root -3, children 5 and 5 solve: Best(take=-3, skip=10)
root -3, children 5 and 5 answer: 10

Each node is visited once, and each call does a fixed amount of work. So the time is proportional to the number of nodes. For n nodes, that is written O(n), which means the work grows in step with n. The brute force method, which tries every group, takes time proportional to 2 to the power n.

Worked example 2: a six-node tree

Here is a tree with six nodes. The root has weight 10. Its left child has weight 1, and the left child has two leaves, 5 and 6. The root's right child has weight 2, and that child has one leaf, 8.

          10
        /    \
       1      2
      / \      \
     5   6      8

Work from the bottom up.

NodeWeighttakeskipHow it was computed
5550leaf, no children
6660leaf, no children
11111take = 1 + 0 + 0; skip = max(5, 0) + max(6, 0) = 11
8880leaf, no children
2228take = 2 + 0; skip = max(8, 0) = 8
10102919take = 10 + 11 + 8; skip = (1's best, 11) + (2's best, 8)

Look at node 1's skip. Its two children, 5 and 6, are not joined to each other. They are siblings. So both can be chosen when node 1 is not, and the skip value is 5 + 6 = 11. Siblings never block each other.

At the root, take is 10 plus the skip values of nodes 1 and 2, which are 11 and 8. That gives 29. The chosen group is {10, 5, 6, 8}. None of these nodes is a parent of another chosen node. Node 10 is joined to 1 and 2, but neither of them is in the group.

The root's skip is 19, from node 1's best (11) and node 2's best (8). The answer is max(29, 19) = 29.

Worked example 3: a deeper tree

This tree has a root of weight 4. Its left child has weight 6, and that child has two leaves of weight 5. The root's right child is a leaf of weight 1.

          4
        /   \
       6     1
      / \
     5   5
Nodetakeskip
5 (left leaf)50
5 (right leaf)50
66 + 0 + 0 = 65 + 5 = 10
1 (leaf)10
4 (root)4 + 10 + 0 = 1410 + 1 = 11

The answer is max(14, 11) = 14, from the group {4, 5, 5}. The node 6 is left out, because it is joined to the root and to both leaves. The two leaves are not joined to the root or to each other, so they can be added together.

Notice that the root's take is only 14 and its skip is 11. Here the root's take wins, so the root is in the best group. The DP finds this by comparing the two values at each step.

Negative weights

Weights can be negative. A node with a negative weight can make a group worse if you choose it. The formulas still work, because take includes the node's own weight, and max picks the better option.

Take a root of weight -3, with two leaves of weight 5 each.

  • take(root) = -3 + 0 + 0 = -3
  • skip(root) = max(5, 0) + max(5, 0) = 10

The answer is max(-3, 10) = 10, from the two leaves alone. The program in the code section prints this case.

An empty group has total 0, so the answer is never below 0. When every node has a negative weight, the answer is 0, from the empty group. The formulas handle this case without a special check, because the base case is 0 and each max compares with the empty choice.

Why one number per node fails

Here is a path with three nodes: A has weight 5, B has weight 10, and C has weight 1. A is the parent of B, and B is the parent of C.

A (5)
  \
   B (10)
     \
      C (1)

The correct answer is 10, from the group {B}. Another group is {A, C}, with total 6. The group {A, B} is not allowed, because A and B are joined.

Now suppose each node sent up only one number, its own weight plus its children's numbers, and never left itself out. For C, the best is 1. For B, the one-number method adds B's weight to C's best, which gives 10 + 1 = 11. That is wrong, because B and C are joined, so they cannot both be chosen. The correct best for B is 10. Then A adds B's best of 11 to its own weight of 5, and gets 16. The group that gives 16 is {A, B, C}, and it is not allowed. The real output of this wrong method is:

Python · editable
path 5-10-1 one-number version: 16

The correct answer for the same path, from the two-value method above, is 10.

The mistake happens because the parent cannot tell whether a child was chosen. With two numbers, the parent knows that it can take its own weight only if each child is skipped. That is exactly the information that the skip value gives.

Common mistakes

Mistake 1: no check for a missing child

Every call to solve reads a child, so the function must stop at a missing node. The code below leaves out the check. It is shown without the run mark because it stops with an error.

from typing import NamedTuple
 
 
class Best(NamedTuple):
    take: int
    skip: int
 
    def best(self):
        return max(self.take, self.skip)
 
 
class Node:
    def __init__(self, weight, left=None, right=None):
        self.weight = weight
        self.left = left
        self.right = right
 
 
def no_null_check(n):
    left = no_null_check(n.left)
    right = no_null_check(n.right)
    return Best(n.weight + left.skip + right.skip, left.best() + right.best())
 
 
path = Node(3, None, Node(4, None, Node(5)))
no_null_check(path)

The output is:

Traceback (most recent call last):
  File "main.py", line 26, in <module>
    no_null_check(path)
    ~~~~~~~~~~~~~^^^^^^
  File "main.py", line 20, in no_null_check
    left = no_null_check(n.left)
  File "main.py", line 20, in no_null_check
    left = no_null_check(n.left)
                         ^^^^^^
AttributeError: 'NoneType' object has no attribute 'left'

The message names the attribute that was read from None. The call on A starts normally and calls itself on A's left child, which is missing (None). That call's first line reads .left from None, so Python stops. The fix is the first if n is None check in solve, which returns Best(0, 0).

Mistake 2: using the best value of the children in take

The take value must add the children's skip values. Using their best values is a common slip, because the names look similar. The code below does this.

Python · editable

On the path A (3), B (4), C (5), this code prints 12, but the correct answer is 8. The result is 12 because the group {A, B, C} is counted. B is joined to A, so it is not allowed. The program prints:

mistake 2 wrong answer: 12 (correct is 8)

Mistake 3: one number per node

The example in the section above shows what a single value gives, with output 16 instead of 10. Keep both values. Use the skip value for a child whenever the parent is taken.

Mistake 4: swapping the two fields

Both values are numbers, so Python does not complain when they are swapped. If the code builds the result as Best(skip, take), the parent then reads the wrong value from each child. The answer for the root can still look reasonable, which makes this mistake hard to see.

Python · editable

The output is:

swapped fields answer: 12 (correct is 8)

The answer 12 is the group {A, B, C}, which is not allowed, because each node in the line is joined to the next one. Name the fields in the constructor call, as Best(take=take, skip=skip), so the order cannot be mixed up.

Mistake 5: deep recursion on a long path

The recursion goes as deep as the tree is tall. A tree that is one long path with 1,000 nodes is more than 1,000 calls deep. Python's default recursion limit, the most calls that can be nested at once, is 1000, so a path of about that size or larger stops with a RecursionError. The code below builds a path of 1,200 nodes, and it is shown without the run mark because it stops with an error.

class Node:
    def __init__(self, weight, left=None, right=None):
        self.weight = weight
        self.left = left
        self.right = right
 
 
def chain(count):
    node = None
    for w in range(count, 0, -1):
        node = Node(1, None, node)
    return node
 
 
def path_best(n):
    if n is None:
        return (0, 0)
    take, skip = path_best(n.right)
    return (n.weight + skip, max(take, skip))
 
 
print(path_best(chain(1200)))

The output is:

Traceback (most recent call last):
  File "main.py", line 22, in <module>
    print(path_best(chain(1200)))
          ~~~~~~~~~^^^^^^^^^^^^^
  File "main.py", line 18, in path_best
    take, skip = path_best(n.right)
                 ~~~~~~~~~^^^^^^^^^
  File "main.py", line 18, in path_best
    take, skip = path_best(n.right)
                 ~~~~~~~~~^^^^^^^^^
  File "main.py", line 18, in path_best
    take, skip = path_best(n.right)
                 ~~~~~~~~~^^^^^^^^^
  [Previous line repeated 996 more times]
RecursionError: maximum recursion depth exceeded

The repeated lines are the same call made again for each node of the path, and the number of repeats depends on the Python version. Two fixes are possible. The first is to raise the limit with sys.setrecursionlimit, which moves the problem to a deeper tree. The second is to rewrite the method with an explicit stack, as the module lesson on iterative traversals shows. The explicit stack is the safer choice for trees that might be very tall, because a higher limit only moves the point where the program fails.

Practice

Exercise 1: a path of three

A root of weight 2 has a child of weight 7. That child has a child of weight 1. Compute take and skip for each node, then give the answer.

Show the answer

The leaf of weight 1 has take 1 and skip 0. The node of weight 7 has take 7 + 0 = 7 and skip max(1, 0) = 1. The root of weight 2 has take 2 + 1 = 3 and skip max(7, 1) = 7. The answer is max(3, 7) = 7. Running solve on this path prints Best(take=3, skip=7), and the answer is 7.

Exercise 2: a root with two children

A root of weight 5 has two leaves, of weight 3 and weight 4. Find the answer and say which group gives it.

Show the answer

The leaves have take 3 and 4, and skip 0 each. The root's take is 5 + 0 + 0 = 5. The root's skip is max(3, 0) + max(4, 0) = 7. The answer is 7, from the group {3, 4}.

Exercise 3: a negative root

A root of weight -1 has two leaves, of weight 2 each. Find the answer.

Show the answer

The leaves have take 2 and skip 0. The root's take is -1 + 0 + 0 = -1. The root's skip is 2 + 2 = 4. The answer is max(-1, 4) = 4. The group is the two leaves, and the root is left out.

Exercise 4: explain the skip value

In the take formula, the parent adds the child's skip value, not the child's take value. Explain in plain words why. Then say what goes wrong if the child's take value is used.

Show the answer

If the parent is in the group, each child is not allowed in the group, because a parent and a child are joined. The skip value is the best total for the child's subtree when the child itself is left out, so it is the value that matches the rule. If the child's take value is used, the group can contain both the parent and the child. The total is then too large, and the answer can be wrong. Mistake 2 in this lesson shows the result on a three-node path.

Exercise 5: recover the group

The DP above returns only the total. Describe in plain words how to find the actual group of nodes after the totals are known. Start from the root.

Show the answer

Keep the take and skip values for every node. At the root, if take is bigger than skip, the root is in the group. Then every child of the root must be left out, so for each child, follow its skip choice. If the root is not in the group, each child can be in or out, and for each child pick the bigger of its take and skip values. Repeat this step down the tree. Each node's choice then follows from its parent's choice, and the nodes you picked form the group.

Summary

  • Each node returns two totals: take, the best total when the node is chosen, and skip, the best total when it is not.
  • A chosen node adds the skip values of its children. A node left out adds the bigger of each child's take and skip values.
  • A missing child returns 0 and 0. The answer for the whole tree is the bigger of the root's two values.
  • Each node is visited once, so the time is O(n) for n nodes.
  • Python integers have no fixed size limit, so the totals cannot overflow. The recursion limit is the limit to watch instead.
Check yourself: Tree DP (6 questions) →Practice this topic: coding problems for Binary Trees →

Stuck? Ask the tutor

Ask about anything in this lesson. Answers are short and use simple words.

Sign in to ask the tutor.

← Zigzag and spiral level order in Python