Tree DP: return two values per node

Why this matters

Imagine a company with an org chart. Each employee has a score, and you want to invite a group of people to a workshop to get the highest total score. There is one rule: you cannot invite a manager and a direct report at the same time, because they would sit in the same meeting and the workshop would not work. Which group gives the best total?

This is a tree problem. Each person is a node. A line between a manager and a report is an edge. The rule says that no two chosen nodes can share an edge. A set of nodes with no edge between any two of them is called an independent set. The problem is to find the independent set with the largest total weight, where the weight of a node is its score. Computer scientists call this the maximum weight independent set problem.

The method in this lesson is called tree dynamic programming, or tree DP. Dynamic programming means solving a big problem by solving its smaller parts once and reusing the answers. On a tree, the smaller parts are subtrees: a node and everything below it. Each node gives its parent a short report about its subtree. The report has two numbers, which is the point of the lesson.

The idea: two numbers per node

Take one node n with a weight. Look at the best total you can get from its subtree, under two conditions.

• with: the best total when n is chosen. If n is chosen, its children cannot be chosen, because each child shares an edge with n. The children's own subtrees are still free, but their roots are out.
• without: the best total when n is not chosen. Its children are now free. Each child may be chosen or not, whichever gives the larger total.

These two numbers are enough to build the answer one level at a time from the bottom up. Here are the two formulas for a node with a left child and a right child:

with(n)    = weight(n) + without(left) + without(right)
without(n) = max(with(left), without(left)) + max(with(right), without(right))

A missing child contributes 0 to both numbers, because there is nothing to add. The answer for the whole tree is max(with(root), without(root)), which is the better of the two cases for the root.

The key point is that each node only needs its children's two numbers. It does not need to know anything else about the tree. That is why the method is efficient: each node is handled once.

Why the formulas look like this

Here is the reasoning in plain words. For with(n), you add the weight of n, and then for each child you can only use the child's "without" value, because the child is forbidden. For without(n), you add, for each child, whichever is better: the child with the child chosen, or the child not chosen. The child is not forbidden in this case, so both options are open. The formula takes the larger of the two.

Worked example 1: a small tree by hand

Use this tree. Each number is a node's weight.

        3
       / \
      2   3
       \   \
        3   1

Work from the bottom up. A leaf (a node with no children) has no children to block, so its with is its weight and its without is 0.

For the root, with is 3 + 3 + 1 = 7. Here the 3 is the without value of the node 2, which is 3, and the 1 is the without value of the right 3, which is 1. The without value of the root is max(2, 3) + max(3, 1) = 3 + 3 = 6. The answer is max(7, 6) = 7.

The chosen set that gives 7 is the root, the bottom 3 under 2, and the 1. No two of them share an edge. Their total is 3 + 3 + 1 = 7.

Worked example 2: where the root must be skipped

Now a tree where the best choice leaves out the root.

        5
       / \
      1   6
     /
    2

The weights are 5 at the root, 1 on the left, 6 on the right, and 2 below the 1.

The answer is max(7, 8) = 8. The best set is 2 and 6. The root 5 is left out, because taking it would block 1 and 6. The node 1 is left out too, because it shares an edge with 2. The node 2 and the node 6 do not share an edge with each other, so both fit.

This example shows why the second number matters. If the code only looked at the root's with value, it would return 7 and miss the answer 8.

Worked example 3: a bigger tree

Now a tree with seven nodes. It is a perfect tree, which means every level is full.

        1
      /   \
     2     3
    / \   / \
   4   5 6   7

Start with the leaves 4, 5, 6 and 7. Each has with equal to its weight and without equal to 0.

Node 2 has children 4 and 5. Its with is 2 + 0 + 0 = 2, because the children's without values are 0. Its without is max(4, 0) + max(5, 0) = 4 + 5 = 9. So node 2 has the pair (2, 9).

Node 3 has children 6 and 7. Its with is 3. Its without is 6 + 7 = 13. So node 3 has the pair (3, 13).

At the root, with is 1 + 9 + 13 = 23. The without value is max(2, 9) + max(3, 13) = 9 + 13 = 22. The answer is max(23, 22) = 23.

The best set is the root, 1, plus all four leaves: 4, 5, 6 and 7. Its total is 1 + 4 + 5 + 6 + 7 = 23. The root is chosen, so its children 2 and 3 are not. The leaves are grandchildren of the root, so they do not share an edge with it. This is a case where taking the root wins, even though the root is small compared with the leaves.

The code

The code builds the tree from level-order values, as in the zigzag lesson. The weights are long long because a sum of many weights can be large, and the same type keeps the arithmetic simple. The Best struct holds the two numbers for one node.

#include <algorithm>
#include <cstdint>
#include <iostream>
#include <memory>
#include <optional>
#include <queue>
#include <vector>
 
struct TreeNode {
    long long val;
    std::unique_ptr<TreeNode> left;
    std::unique_ptr<TreeNode> right;
    explicit TreeNode(long long v) : val(v) {}
};
 
using NodePtr = std::unique_ptr<TreeNode>;
 
// Builds a tree from level-order values. std::nullopt marks a missing child.
NodePtr build(const std::vector<std::optional<long long>>& vals) {
    if (vals.empty() || !vals[0]) return nullptr;
    NodePtr root = std::make_unique<TreeNode>(*vals[0]);
    std::queue<TreeNode*> q;
    q.push(root.get());
    std::size_t i = 1;
    while (!q.empty() && i < vals.size()) {
        TreeNode* cur = q.front();
        q.pop();
        if (i < vals.size() && vals[i]) {
            cur->left = std::make_unique<TreeNode>(*vals[i]);
            q.push(cur->left.get());
        }
        ++i;
        if (i < vals.size() && vals[i]) {
            cur->right = std::make_unique<TreeNode>(*vals[i]);
            q.push(cur->right.get());
        }
        ++i;
    }
    return root;
}
 
struct Best {
    long long with;     // best total in this subtree when this node is chosen
    long long without;  // best total in this subtree when this node is not chosen
};
 
Best solve(const TreeNode* n) {
    if (!n) return {0, 0};                       // an empty subtree adds nothing
    Best l = solve(n->left.get());
    Best r = solve(n->right.get());
    long long with = n->val + l.without + r.without;
    long long without = std::max(l.with, l.without) + std::max(r.with, r.without);
    return {with, without};
}
 
long long max_independent_set(const TreeNode* root) {
    Best b = solve(root);
    return std::max(b.with, b.without);          // the root may or may not be chosen
}

The function solve does the whole method. It calls itself (this is recursion, where a function calls itself) on the left and right child first, so the children's values are ready before the parent needs them. This is the bottom-up order. The base case if (!n) return {0, 0}; handles a missing child. The two formulas from the idea section are the two lines that compute with and without.

The function max_independent_set takes the better of the root's two values, because the root may or may not be part of the best set.

Checking the code

A good habit is to check a new method against a brute force method (one that tries every possibility) that is certainly right. For a small tree, the slow method tries every subset of nodes, keeps only the subsets where no two chosen nodes share an edge, and takes the largest total. With 12 nodes there are 4096 subsets, which runs in a fraction of a second.

void collect(const TreeNode* n, std::vector<const TreeNode*>& nodes) {
    if (!n) return;
    nodes.push_back(n);
    collect(n->left.get(), nodes);
    collect(n->right.get(), nodes);
}
 
// Slow but certain: try every subset of nodes.
long long brute(const TreeNode* root) {
    std::vector<const TreeNode*> nodes;
    collect(root, nodes);
    std::size_t m = nodes.size();
    long long best = 0;
    for (std::uint32_t mask = 0; mask < (1u << m); ++mask) {
        bool ok = true;
        long long sum = 0;
        for (std::size_t i = 0; i < m && ok; ++i) {
            if (!(mask & (1u << i))) continue;
            sum += nodes[i]->val;
            for (std::size_t j = 0; j < m; ++j) {
                if (!(mask & (1u << j))) continue;
                const TreeNode* a = nodes[i];
                if (a->left.get() == nodes[j] || a->right.get() == nodes[j]) { ok = false; break; }
            }
        }
        if (ok) best = std::max(best, sum);
    }
    return best;
}

The check is in the inner loop: a chosen node and another chosen node are not allowed to be parent and child.

The main function below prints the examples from this lesson, along with the empty tree and a tree with a negative weight.

int main() {
    using std::nullopt;
    NodePtr t1 = build({3, 2, 3, nullopt, 3, nullopt, 1});
    std::cout << "example 1: " << max_independent_set(t1.get())
              << "  brute force: " << brute(t1.get()) << "\n";
 
    NodePtr t2 = build({5, 1, 6, 2});
    std::cout << "example 2: " << max_independent_set(t2.get())
              << "  brute force: " << brute(t2.get()) << "\n";
 
    NodePtr t3 = build({1, 2, 3, 4, 5, 6, 7});
    std::cout << "example 3: " << max_independent_set(t3.get())
              << "  brute force: " << brute(t3.get()) << "\n";
 
    NodePtr t4 = build({});
    std::cout << "empty tree: " << max_independent_set(t4.get()) << "\n";
 
    NodePtr t5 = build({-4, 5, 2});
    std::cout << "negative root: " << max_independent_set(t5.get())
              << "  brute force: " << brute(t5.get()) << "\n";
    return 0;
}

Compile and run it with:

g++ -std=c++20 -Wall -Wextra -g mis.cpp -o mis && ./mis

The output is:

example 1: 7  brute force: 7
example 2: 8  brute force: 8
example 3: 23  brute force: 23
empty tree: 0
negative root: 7  brute force: 7

The two methods agree on every example. The negative-weight case works without a special rule. The root has weight -4, so the best set leaves it out, and the answer comes from the two children, 5 and 2. The set can be empty, so the code never has to force a node into the answer.

To test on more trees, a random test builds 3000 trees with between 1 and 12 nodes, each with a weight from -5 to 15, and compares the two methods. In one run of that test, the output was:

random trees checked: 3000 mismatches: 0

A test like this does not prove the method is right for every tree, but it catches most mistakes in the formulas.

Common mistakes

Adding every weight

A common first idea is to add the weight of every node. For example 1, that gives 3 + 2 + 3 + 3 + 1 = 12. The rule allows no two chosen nodes to share an edge, so the answer must be smaller: 7. A program that sums every node prints 12 and reports no error. The wrong number is the only sign of the bug.

Using only the without values of the children

A common bug is to write without = l.without + r.without. This ignores the case where a child is chosen, so a node that is not chosen gets no credit for its children's best choice. Build this version with the same tree code and run it on example 2. It prints 5. The correct answer is 8, from the set 2 and 6. The bug hides the 2 because node 1 never gets credit for it, and the root then sees a total of 0 from its left side.

Returning only the with value

The root has two cases, and the answer needs the better one. A version that returns solve(root).with gives 7 on example 2, but the correct answer is 8. Always take the maximum of the root's two numbers.

Forgetting the empty-child values

If a missing child returned values that are too large, the answer would include nodes that do not exist. The base case {0, 0} is the right value: an empty subtree adds nothing, and it blocks nothing.

Practice

Exercise 1. Use the method by hand on the chain of three nodes with weights 3, 2 and 3, from the top: the root is 3, its child is 2, and the child of 2 is 3. Give the answer.

Show the answer

The bottom 3 has the pair (3, 0). The 2 has with = 2 + 0 = 2 and without = max(3, 0) = 3, so the pair is (2, 3). The root has with = 3 + 3 = 6 and without = max(2, 3) = 3. The answer is max(6, 3) = 6. The chosen set is the root and the bottom 3, which are not neighbours. The brute force check gives 6 too.

Exercise 2. The root has weight 10 and two children, each with weight 1 and no children. Give the answer and explain why the children are not chosen.

Show the answer

Each child has the pair (1, 0). The root has with = 10 + 0 + 0 = 10. Its without is max(1, 0) + max(1, 0) = 2. The answer is max(10, 2) = 10. Choosing the root blocks both children, since each shares an edge with it. The root alone is better than the two children together, because the children together give only 2.

Exercise 3. The tree has a root of -4, with children 5 and 2. Explain why the answer is 7 and not -4.

Show the answer

The children have the pairs (5, 0) and (2, 0). The root has with = -4 + 0 + 0 = -4, which is bad. Its without value is max(5, 0) + max(2, 0) = 7. The answer is max(-4, 7) = 7. The root is left out, which is allowed, so the set wins.

Exercise 4. Explain in plain words why an empty child returns the pair (0, 0). What would go wrong if it returned (-1, -1) instead?

Show the answer

An empty child has no weight, so it should not change any total. A value of 0 adds nothing. If an empty child returned -1, each parent would lose 1 from both of its numbers for every missing child, so the answer would come out too small even when every weight is correct.

Exercise 5. For the tree in example 1, name the set of nodes that gives the answer 7.

Show the answer

The set is the root (3), the bottom 3 under the left 2 (3), and the 1 under the right 3 (1). Their total is 7. The root is adjacent only to its two children, which are the 2 and the right 3, and neither of them is chosen. The bottom 3 and the 1 are grandchildren of the root, so they share no edge with it or with each other.

Summary

• Each node reports two numbers: the best total in its subtree with the node chosen, and the best total with the node not chosen.
• A chosen node forces its children to be unchosen, so with uses the children's without values.
• A node that is not chosen lets each child be chosen or not, so without takes the larger value for each child.
• The answer for the whole tree is the larger of the root's two numbers.
• A missing child contributes {0, 0}, and a brute force check on small random trees is a good way to test the formulas.