Skip to content

Tree Coloring

From Two States to Many

Every problem in this chapter so far had two states per node: take or skip, in cover or not, matched or unmatched. Coloring generalizes this to \(k\) states.

A proper \(k\)-coloring assigns one of \(k\) colors to each node so that no two adjacent nodes share the same color. The question: how many valid \(k\)-colorings does a tree have?

Trees Are Always 2-Colorable

First, the existence question. Can every tree be 2-colored? Yes — trees are bipartite. Run BFS from the root, alternating colors at each level. Since trees have no cycles, you'll never hit a conflict.

This is a well-known fact, but it connects directly to this chapter's theme: bipartiteness is the "coloring" version of the include-exclude structure. At each node, you choose one of two states. Adjacency constraints restrict which states neighbors can share.

Each non-root node has k-1 color choices since it must differ from its parent

Counting: The Closed-Form Formula

How many valid \(k\)-colorings does a tree with \(n\) nodes have?

Root the tree at any node. The root has \(k\) choices. Every other node must differ from its parent, so it has \(k - 1\) choices. Since the tree structure means each non-root node has exactly one parent, and subtrees are independent:

\[\text{count} = k \cdot (k-1)^{n-1}\]

That's it. No DP needed for this version.

Example: The path 1-2-3 with \(k = 3\) colors. Root at node 1: 3 choices. Node 2: must differ from 1, so 2 choices. Node 3: must differ from 2, so 2 choices. Total: \(3 \times 2 \times 2 = 12\).

Verification for \(k = 2\): A path of 3 nodes with 2 colors. Root gets 2 choices. Each subsequent node gets 1 choice. Total: \(2 \times 1 \times 1 = 2\). These are: (R, B, R) and (B, R, B). Correct.

Why the formula works: In a tree, the constraint graph is a tree itself. Each node's color is constrained only by its parent. No cycles means no additional constraints beyond parent-child.

When the Formula Isn't Enough

The formula \(k \cdot (k-1)^{n-1}\) assumes every "adjacent nodes must differ" constraint. But what if you have mixed constraints?

  • Some edges require "must differ" (standard coloring)
  • Some edges require "must agree" (same color)
  • Some nodes are pre-colored

Now you need actual DP. And this is where coloring connects to the include-exclude framework.

The DP Approach

Define \(\text{dp}[\text{node}][c]\) = number of valid colorings of this node's subtree when this node is assigned color \(c\).

Leaf node: \(\text{dp}[\text{leaf}][c] = 1\) for every color \(c\).

Internal node with "must differ" constraint: For each color \(c\) assigned to this node, each child can take any color except \(c\):

\[\text{dp}[\text{node}][c] = \prod_{\text{child}} \left( \sum_{c' \neq c} \text{dp}[\text{child}][c'] \right)\]

The total for the tree is \(\sum_{c} \text{dp}[\text{root}][c]\).

Full Trace

        1
       / \
      2    3
     / \
    4   5

Using \(k = 3\) colors (call them 0, 1, 2). Process bottom-up.

Leaves (nodes 3, 4, 5): dp = [1, 1, 1] for each.

Node 2 (children: 4, 5):

For each color \(c\) of node 2, child 4 can be any color except \(c\), and child 5 can be any color except \(c\).

Color of node 2 Valid colors for child 4 Count from child 4 Valid colors for child 5 Count from child 5 dp[2][c]
0 {1, 2} 1 + 1 = 2 {1, 2} 2 2 * 2 = 4
1 {0, 2} 1 + 1 = 2 {0, 2} 2 4
2 {0, 1} 1 + 1 = 2 {0, 1} 2 4

So dp[2] = [4, 4, 4].

Node 1 (children: 2, 3):

Color of node 1 Valid from child 2 Valid from child 3 dp[1][c]
0 dp[2][1] + dp[2][2] = 4 + 4 = 8 dp[3][1] + dp[3][2] = 1 + 1 = 2 8 * 2 = 16
1 dp[2][0] + dp[2][2] = 4 + 4 = 8 dp[3][0] + dp[3][2] = 1 + 1 = 2 16
2 dp[2][0] + dp[2][1] = 4 + 4 = 8 dp[3][0] + dp[3][1] = 1 + 1 = 2 16

Total: \(16 + 16 + 16 = 48\).

Formula check: \(k \cdot (k-1)^{n-1} = 3 \cdot 2^4 = 3 \cdot 16 = 48\). Matches.

Constraint Propagation

The DP above is doing constraint propagation — the color choice at a parent restricts what's available at each child. In the standard "must differ" case, the restriction is symmetric and leads to the clean formula.

But consider a mixed scenario: edge (1,2) requires "must agree" while all other edges require "must differ." Now node 2 MUST have the same color as node 1. The formula breaks — you need the DP.

This is the same structure as include-exclude, just with more states:

Framework States per node Constraint
Independent Set {skip, take} take → children skip
Vertex Cover {out, in} out → children in
\(k\)-Coloring {color 0, color 1, ..., color \(k{-}1\)} color \(c\) → children not \(c\)

Coloring IS include-exclude with \(k\) states instead of 2. The recurrence has the same product-over-children shape. The only difference is the number of options at each node.

The Code

For the general DP (works with any constraint pattern):

vector<long long> solve(TreeNode* root, int numColors) {
    vector<long long> colorCount(numColors, 1);
    if (!root) return vector<long long>(numColors, 1);

    auto leftCounts = root->left ? solve(root->left, numColors)
                                 : vector<long long>(numColors, 1);
    auto rightCounts = root->right ? solve(root->right, numColors)
                                   : vector<long long>(numColors, 1);

    long long totalLeft = 0;
    long long totalRight = 0;
    for (int c = 0; c < numColors; c++) {
        totalLeft += leftCounts[c];
        totalRight += rightCounts[c];
    }

    for (int c = 0; c < numColors; c++) {
        long long validLeft = totalLeft - leftCounts[c];
        long long validRight = totalRight - rightCounts[c];
        colorCount[c] = validLeft * validRight;
    }

    return colorCount;
}

long long countColorings(TreeNode* root, int numColors) {
    if (!root) return 0;
    auto counts = solve(root, numColors);
    long long total = 0;
    for (long long c : counts) total += c;
    return total;
}

Complexity: \(O(n \cdot k)\) time, \(O(h \cdot k)\) space.

The Chromatic Polynomial

What you just computed is actually a famous object: the chromatic polynomial \(P(T, k)\) of a tree \(T\). For any tree with \(n\) nodes, \(P(T, k) = k \cdot (k-1)^{n-1}\). This polynomial tells you the number of valid \(k\)-colorings for any positive integer \(k\).

Plug in \(k = 1\): \(1 \cdot 0^{n-1} = 0\) (can't color any tree with one color if \(n > 1\)). Plug in \(k = 2\): \(2 \cdot 1^{n-1} = 2\) (exactly two 2-colorings for any tree). These match intuition.

For general graphs, the chromatic polynomial is much harder to compute. But for trees, the lack of cycles makes everything factor cleanly.

Try It

Input: (5-node tree, 3 colors)
Output: 48

Predict before running: How many valid 2-colorings does a 5-node tree have? \(2 \cdot 1^4 = 2\). Always just 2 — swap the two colors.

Challenge: A tree has 10 nodes. How many valid 4-colorings? \(4 \cdot 3^9 = 4 \cdot 19683 = 78732\).

Challenge 2: If one node is pre-colored (say the root must be red), how does the formula change? The root no longer has \(k\) choices — it has 1. So the count becomes \((k-1)^{n-1}\).

Edge Cases to Watch For

  • \(k = 1\) with \(n > 1\): The formula gives \(1 \times 0^{n-1} = 0\), which is correct (two adjacent nodes can't share one color). But if your code computes \(0^0\) for the \(n=1\) subcase, some languages return unexpected results — guard this explicitly.
  • Pre-colored node: If one node is fixed to a specific color, the root has 1 choice instead of \(k\). The count becomes \((k-1)^{n-1}\). Forgetting to adjust the root factor is a common bug when constraints add fixed colors.
  • Integer overflow in the formula: \(k \cdot (k-1)^{n-1}\) grows exponentially. For \(n = 20\) and \(k = 10\), the result exceeds \(10^{19}\), overflowing long long. Use modular arithmetic if the problem requires it.

Problems

Problem Link Difficulty
LC 886 — Possible Bipartition leetcode.com/problems/possible-bipartition Medium
LC 785 — Is Graph Bipartite? leetcode.com/problems/is-graph-bipartite Medium
CF 902B — Coloring a Tree codeforces.com/problemset/problem/902/B Easy