AtCoder - Subtree

Solving for One Root

Let's consider a simpler problem - assuming that node 1 is painted black, how many ways can we paint the tree?

First, root the tree at node 1.

Let $dp[i]$ be the number of ways that we can paint the subtree of node $i$ if we know that node $i$ is painted black. (1 if node $i$ is a leaf)

For each child $c$ of $i$, there are $dp[c] + 1$ ways to paint its subtree if $i$ is painted black.

This means that we have the recurrence

The answer to the simpler problem is thus $dp[1]$.

Finding all $dp[i]$ takes $O(N)$ time.

Solving for all roots

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

Let $dp2[i]$ be the number of ways to colour the tree if we remove node $i$'s subtree excluding node $i$ from the graph and we know that node $i$ is painted black. (1 if node $i$ is the root)

The number of ways to paint the tree if we know node $i$ is black is simply $dp[i] \cdot dp2[i]$.

How can we find $dp2[i]$ efficiently though?

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

Since $M$ 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]$. (Since we can't find modular inverses easily.)

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

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

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

Code

down corresponds to $dp$ and up corresponds to $dp2$.

/**
 * Description: modular arithmetic operations
 */

template<int MOD, 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
    mint() { v = 0; }
    mint(ll _v) { v = (-MOD < _v && _v < MOD) ? _v : _v % MOD;