AtCoder DP Contest - Subtree

Authors: Benjamin Qi, Andi Qu

Solving For One Root

Let's consider a simpler problem:

Assuming that node 11 is painted black, how many ways can we paint the tree?

First, root the tree at node 11. Let dp[i]dp[i] be the number of ways that we can paint the subtree of node ii such that either node ii is colored black, or no nodes are colored black. Note that if ii is a leaf, then dp[i]=2dp[i]=2 (we choose to color node ii black or not).

For each child cc of ii, there are dp[c]dp[c] ways to paint its subtree if ii is painted black. This means that we have the recurrence

dp[i]=1+cChildren of idp[c]dp[i]=1+\prod_{c \in \text{Children of } i} dp[c]

where the product corresponds to painting node ii black and the 11 corresponds to painting node ii white.

The answer to the simpler problem is thus dp[1]1dp[1]-1. Finding all dp[i]dp[i] takes O(N)\mathcal{O}(N) time.

Solving For All Roots

First, root the tree arbitrarily and do a DFS to find all dp[i]dp[i].

Let dp2[i]dp2[i] be the number of ways to colour the tree if we remove node ii's subtree such that either the parent of ii is black, or no nodes are colored black. Note that dp2[1]=1dp2[1]=1.

The number of ways to paint the tree if we know node ii is black is simply (dp[i]1)dp2[i](dp[i]-1)\cdot dp2[i]. How can we find dp2[i]dp2[i] efficiently though?

The basic recurrence for computing dp2[i]dp2[i] is

dp2[i]=1+dp2[Parent of i]sSiblings of idp[s]dp2[i] = 1+dp2[\text{Parent of } i] \cdot \prod_{s \in \text{Siblings of } i} dp[s]

where the product corresponds to painting the parent of ii black and the 11 corresponds to painting the parent of ii white.

Since MM isn't guaranteed to be prime though, we can't just find the product of each node's children and divide that product by each of its children's dp[i]dp[i] (since we can't find modular inverses easily).

However, notice how if node ii is the kk-th child of its parent, then we can use prefix and suffix products to compute

sSiblings of idp[s]\prod_{s \in \text{Siblings of } i}dp[s]

without using division. (i.e. We find the product of dp[s]dp[s] for the first to the (k1)(k - 1)-th child of ii's parent, the product of dp[s]dp[s] for the (k+1)(k + 1)-th to the last child of ii's parent, and then multiply those together.)

Finding all dp2[i]dp2[i] takes O(N)\mathcal{O}(N) time using a DFS, so the total complexity of this algorithm is thus O(N)\mathcal{O}(N).


down corresponds to dpdp and up corresponds to dp2dp2. The code uses the exact same recurrences mentioned above.

Code Snippet: Benq Template (Click to expand)
* Description: modular arithmetic operations
template<int RT> struct mint {
// static const int mod = MOD;
static constexpr mint rt() { return RT; } // primitive root for FFT
int v; explicit operator int() const { return v; } // explicit -> don't silently convert to int

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!