Árbol virtual

A menudo, al ejecutar operaciones o programación dinámica sobre árboles, solo necesitamos llevar la cuenta de unos pocos nodos clave y sus relaciones. Si podemos aislar los pocos nodos relevantes que mantienen intacta la estructura del árbol, conservando a la vez las relaciones entre los nodos clave, podemos reducir de forma significativa nuestra complejidad temporal.

Ejemplo - Leaf Color

Solución naive

Podemos sumar la cantidad de conjuntos de vértices $T$ que satisfacen la condición de que los vértices de grado $1$ tienen color $c$ para todo $c \in [1,N]$ para obtener una respuesta.

Definimos $\texttt{dp}[v]$ como la cantidad de subgrafos inducidos que son árboles enraizados en $v$ tales que todas las hojas tienen color $c$. Al hacer las transiciones, inicializamos $\texttt{dp}[v] = 1$. Las transiciones son las siguientes:

• Para cada hijo $u$ de $v$, podemos o bien adjuntar uno de los subgrafos inducidos que son árboles que contienen a $u$, o ignorar $u$ por completo:

  $$\texttt{dp}[v] = \texttt{dp}[v] \cdot (1 + \texttt{dp}[u])$$

• Sin embargo, si $a_v \neq c$, entonces tenemos que quitar el caso en que el subgrafo inducido es solo $v$ mismo, porque viola la condición:

  $$\texttt{dp}[v] = \texttt{dp}[v] - 1$$

Para extraer nuestra respuesta, podemos iterar sobre cada raíz posible $r$ de nuestro subgrafo inducido.

• Si $a_r =c$, entonces podemos sumar $\texttt{dp}[v]$ directamente a nuestra respuesta.
• En caso contrario, hay que asegurarse de que $r$ no tenga grado $1$ en ninguno de los subgrafos inducidos que sumamos. Así, restamos todos los casos en que $r$ está unido a exactamente uno de sus hijos.

Así, nuestra respuesta será

La complejidad temporal de este enfoque es $\mathcal{O}(n^2)$, porque para cada color iteramos sobre todos los $n$ nodos.

Aquí hay un código (TLE) que demuestra este enfoque:

#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
 
constexpr int MAX_N = 2e5 + 1;
constexpr int MOD = 998244353;
 
int n, a[MAX_N];
ll dp[MAX_N], ans;
vector<int> adj[MAX_N];
 
int dfs_dp(int v, int p, int c) {
	dp[v] = 1;
	for (int u : adj[v]) {
		if (u == p) continue;
		dp[v] *= (dfs_dp(u, v, c) + 1);
		dp[v] %= MOD;
	}
	if (a[v] != c) dp[v]--;
 
	ans += dp[v];
	ans %= MOD;
 
	if (a[v] != c) {
		for (int u : adj[v]) {
			if (u == p) continue;
			ans += MOD - dp[u];
			ans %= MOD;
		}
	}
 
	return dp[v];
}
 
int main() {
	cin >> n;
 
	for (int i = 0; i < n; i++) { cin >> a[i]; }
 
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		u--;
		v--;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
 
	for (int col = 1; col <= n; col++) { dfs_dp(0, 0, col); }
 
	cout << ans;
}

Optimización con árbol virtual

Para acelerar nuestro enfoque de programación dinámica, podemos observar que nuestras transiciones de programación dinámica son relativamente simples cuando involucran ciertos tipos de nodos.

El caso más simple es cuando un nodo $v$ no tiene nodos de color $c$ en su subárbol. Aquí, simplemente concluimos $\texttt{dp[v]}=0$, porque es imposible construir un subgrafo inducido en el subárbol de $v$ con hojas de color $c$.

Otro caso de transiciones simples es cuando solo un hijo $u$ de $v$ tiene algún nodo de color $c$. Entonces, a partir de las transiciones de DP ya establecidas, observamos que $\texttt{dp}[v]$ es equivalente a $\texttt{dp}[u]$.

Sea $S$ el conjunto de nodos excluyendo estos dos casos. Más formalmente, si hacemos $S_c = \{ v \in V \mid a_v = c \}$, entonces

Definimos $S$ como el árbol virtual o árbol auxiliar de $S_c$. También podemos definir $p'_v$, el padre virtual de $v$, como el nodo más bajo en $S$ tal que $p'_v \neq v$ y $\text{LCA}(p'_v,v) = p'_v$, y $v$ como un hijo virtual de $p'_v$.

Aquí hay un ejemplo de un árbol virtual, con los nodos de $S_c$ coloreados de púrpura, y el resto del árbol virtual en azul:

A partir de aquí, podemos calcular nuestros valores de DP de forma similar a antes, excepto que solo procesamos nodos en $S$, y consideramos hijos virtuales en lugar de hijos directos.

Observar que al calcular nuestra respuesta, también solo necesitamos considerar nodos en $S$. Los nodos que no están en $S$ no pueden ser la raíz porque solo uno de sus hijos tiene nodos de color $c$ en su subárbol, lo que significa que terminaríamos con un nodo de grado $1$ que no es de color $c$.

Así, nuestra complejidad temporal será $\mathcal{O}(\sum_c |S|)$, ya que iteramos sobre el conjunto $S$ para cada color distinto.

En la siguiente sección, demostramos que $|S|$ está acotado por $\mathcal{O}(S_c)$, y también mostramos cómo construir un árbol virtual.

Construcción del árbol virtual

Un método bien conocido para construir el árbol virtual del conjunto $S_0$ es el siguiente:

• Ordenar los puntos de $S_0$ por orden DFS, y agregarlos a $S$
• Calcular el LCA de cualesquiera dos puntos clave adyacentes en $S_0$, y agregarlo a $S$
• Construir el árbol virtual según la relación ancestro-descendente del árbol original

Para una demostración informal de correctitud, considerar lo siguiente:

Observar que un vértice $v \in S$ debe o bien satisfacer $v \in S_0$ o existir al menos dos hijos distintos de $v$ con puntos clave en sus subárboles.

El primer caso se resuelve de inmediato en el paso 1. En el segundo caso, como propiedad del orden DFS, deben existir puntos clave $a,b$ en el subárbol de $v$ tales que $a,b$ están en subárboles distintos y $a,b$ son adyacentes en el orden DFS cuando solo se consideran los puntos clave. Como $v$ es $\text{LCA}(a,b)$, nuestra construcción necesariamente se cumple. Un corolario de nuestra construcción es que $|S|$ está de hecho acotado por $\mathcal{O}(|S_0|)$: de hecho, una cota ajustada es $|S| \leq 2 \cdot |S_0|-1$.

Implementación

Abajo hay una implementación para la tarea, que también contiene código general para construir árboles virtuales. Como $\sum_c |S|$ es $\mathcal{O(n)}$, y ordenamos por orden DFS, nuestra complejidad temporal final es $\mathcal{O}(n\log n)$.

#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
 
constexpr int MAX_N = 2e5 + 1;
constexpr int LG = 18;
constexpr int MOD = 998244353;
 
int n, a[MAX_N];
ll dp[MAX_N], ans;
int tin[MAX_N], tout[MAX_N], d[MAX_N], lift[MAX_N][LG], timer;
vector<int> adj[MAX_N], vadj[MAX_N], at_a[MAX_N];
 
void dfs(int v, int p) {
	at_a[a[v]].push_back(v);
	tin[v] = timer++;
 
	lift[v][0] = p;
	for (int i = 1; i < LG; i++) { lift[v][i] = lift[lift[v][i - 1]][i - 1]; }
 
	for (int u : adj[v]) {
		if (u == p) { continue; }
		dfs(u, v);
	}
 
	tout[v] = timer++;
}
 
int is_ancestor(int u, int v) { return tin[u] <= tin[v] && tout[v] <= tout[u]; }
 
int lca(int u, int v) {
	if (is_ancestor(u, v)) { return u; }
	if (is_ancestor(v, u)) { return v; }
 
	for (int i = LG - 1; i >= 0; i--) {
		if (!is_ancestor(lift[u][i], v)) { u = lift[u][i]; }
	}
	return lift[u][0];
}
 
bool sort_tin(const int &a, const int &b) { return tin[a] < tin[b]; }
 
vector<int> vtree(const vector<int> &key) {
	if (key.empty()) return {};
 
	vector<int> res = key;
	sort(res.begin(), res.end(), sort_tin);
 
	for (int i = 1; i < (int)key.size(); i++) {
		res.push_back(lca(res[i - 1], res[i]));
	}
 
	sort(res.begin(), res.end(), sort_tin);
	res.erase(unique(res.begin(), res.end()), res.end());
 
	for (int v : res) { vadj[v].clear(); }
 
	for (int i = 1; i < (int)res.size(); i++) {
		vadj[lca(res[i - 1], res[i])].push_back(res[i]);
	}
 
	return res;
}
 
int main() {
	cin >> n;
 
	for (int i = 0; i < n; i++) { cin >> a[i]; }
 
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		u--;
		v--;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
 
	dfs(0, 0);
 
	for (int col = 1; col <= n; col++) {
		vector<int> vt = vtree(at_a[col]);
		reverse(vt.begin(), vt.end());
 
		for (int v : vt) {
			dp[v] = 1;
			for (int u : vadj[v]) {
				dp[v] *= (dp[u] + 1);
				dp[v] %= MOD;
			}
			if (a[v] != col) { dp[v]--; }
 
			ans += dp[v];
			ans %= MOD;
			if (a[v] != col) {
				for (int u : vadj[v]) {
					ans += MOD - dp[u];
					ans %= MOD;
				}
			}
		}
	}
 
	cout << ans << '\n';
}

Problemas