Subtree

Resolver para una raíz

Consideremos un problema más simple:

Suponiendo que el nodo $1$ está pintado de negro, ¿de cuántas formas podemos pintar el árbol?

Primero, enraizamos el árbol en el nodo $1$ . Sea $dp[i]$ la cantidad de formas en que podemos pintar el subárbol del nodo $i$ de modo que o bien el nodo $i$ esté coloreado de negro, o bien ningún nodo esté coloreado de negro. Notamos que si $i$ es una hoja, entonces $dp[i]=2$ (elegimos colorear el nodo $i$ de negro o no).

Para cada hijo $c$ de $i$ , hay $dp[c]$ formas de pintar su subárbol si $i$ está pintado de negro. Esto significa que tenemos la recurrencia

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

donde el producto corresponde a pintar el nodo $i$ de negro y el $1$ corresponde a pintar el nodo $i$ de blanco.

La respuesta al problema más simple es entonces $dp[1]-1$ . Hallar todos los $dp[i]$ toma tiempo $\mathcal{O}(N)$ .

Resolver para todas las raíces

Primero, enraizamos el árbol de forma arbitraria y hacemos un DFS para hallar todos los $dp[i]$ .

Sea $dp2[i]$ la cantidad de formas de colorear el árbol si quitamos el subárbol del nodo $i$ de modo que o bien el padre de $i$ es negro, o bien ningún nodo está coloreado de negro. Notamos que $dp2[1]=1$ .

La cantidad de formas de pintar el árbol si sabemos que el nodo $i$ es negro es simplemente $(dp[i]-1)\cdot dp2[i]$ . ¿Cómo podemos hallar $dp2[i]$ de forma eficiente?

La recurrencia básica para calcular $dp2[i]$ es

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

donde el producto corresponde a pintar el padre de $i$ de negro y el $1$ corresponde a pintar el padre de $i$ de blanco.

Como no está garantizado que $M$ sea primo, no podemos simplemente hallar el producto de los hijos de cada nodo y dividir ese producto por el $dp[i]$ de cada hijo (porque no podemos hallar inversos modulares con facilidad).

Sin embargo, notamos que si el nodo $i$ es el $k$ -ésimo hijo de su padre, entonces podemos usar productos de prefijos y sufijos para calcular

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

sin usar división. (Es decir, hallamos el producto de $dp[s]$ para el primero hasta el $(k - 1)$ -ésimo hijo del padre de $i$ , el producto de $dp[s]$ para el $(k + 1)$ -ésimo hasta el último hijo del padre de $i$ , y luego multiplicamos esos juntos.)

Hallar todos los $dp2[i]$ toma tiempo $\mathcal{O}(N)$ usando un DFS, así que la complejidad total de este algoritmo es $\mathcal{O}(N)$ .

Implementación

down corresponde a $dp$ y up corresponde a $dp2$ . El código usa las mismas recurrencias mencionadas arriba.

C++


// CodeSnip{Benq Template}
 
/**
 * 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
	mint() { v = 0; }
	mint(ll _v) {
		v = int((-MOD < _v && _v < MOD) ? _v : _v % MOD);
		if (v < 0) v += MOD;
	}
	friend bool operator==(const mint &a, const mint &b) { return a.v == b.v; }
	friend bool operator!=(const mint &a, const mint &b) { return !(a == b); }
	friend bool operator<(const mint &a, const mint &b) { return a.v < b.v; }
	friend void re(mint &a) {
		ll x;
		re(x);
		a = mint(x);
	}
	friend str ts(mint a) { return ts(a.v); }
 
	mint &operator+=(const mint &m) {
		if ((v += m.v) >= MOD) v -= MOD;
		return *this;
	}
	mint &operator-=(const mint &m) {
		if ((v -= m.v) < 0) v += MOD;
		return *this;
	}
	mint &operator*=(const mint &m) {
		v = (ll)v * m.v % MOD;
		return *this;
	}
	mint &operator/=(const mint &m) { return (*this) *= inv(m); }
	friend mint pow(mint a, ll p) {
		mint ans = 1;
		assert(p >= 0);
		for (; p; p /= 2, a *= a)
			if (p & 1) ans *= a;
		return ans;
	}
	friend mint inv(const mint &a) {
		assert(a.v != 0);
		return pow(a, MOD - 2);
	}
 
	mint operator-() const { return mint(-v); }
	mint &operator++() { return *this += 1; }
	mint &operator--() { return *this -= 1; }
	friend mint operator+(mint a, const mint &b) { return a += b; }
	friend mint operator-(mint a, const mint &b) { return a -= b; }
	friend mint operator*(mint a, const mint &b) { return a *= b; }
	friend mint operator/(mint a, const mint &b) { return a /= b; }
};
 
typedef mint<5> mi;
 
template <int SZ> struct SubtreeDP {
	int par[SZ];
	vi adj[SZ];
	void ae(int a, int b) { adj[a].pb(b), adj[b].pb(a); }
	struct T {
		mi v = 1;
		T &operator+=(const T &b) {
			v *= b.v;
			return *this;
		}
		void tran() { ++v; }
	};
	T up[SZ], down[SZ];
	void dfs(int x) {
		trav(t, adj[x]) if (t != par[x]) {
			par[t] = x;
			dfs(t);
			down[x] += down[t];
		}
		down[x].tran();
	}
	void dfs2(int x) {
		{
			T pre = up[x];  // deal with prefixes
			F0R(i, sz(adj[x])) {
				int c = adj[x][i];
				if (c == par[x]) continue;
				up[c] += pre;
				pre += down[c];
			}
		}
		{
			T pre;  // deal with suffixes
			R0F(i, sz(adj[x])) {
				int c = adj[x][i];
				if (c == par[x]) continue;
				up[c] += pre;
				pre += down[c];
			}
		}
		F0R(i, sz(adj[x])) {
			int c = adj[x][i];
			if (c == par[x]) continue;
			up[c].tran();
			dfs2(c);
		}
	}
	// T getSub(int x, int y) { return par[x] == y ? down[x] : up[y]; }
	// get subtree of x excluding y
	void init(int n) {
		par[1] = 0;
		dfs(1);
		dfs2(1);
		FOR(i, 1, n + 1) ps((down[i].v - 1) * up[i].v);
		// FOR(i,1,n+1) { alternative method
		//     T p = T(); trav(t,adj[i]) p += getSub(t,i);
		//     ps(p.v);
		// }
	}
};
 
int main() {
	setIO();
	int n;
	re(n, MOD);
	SubtreeDP<MX> S;
	F0R(i, n - 1) {
		int a, b;
		re(a, b);
		S.ae(a, b);
	}
	S.init(n);
}