Fixed-Length Paths I

Complejidad temporal: $\mathcal O(N \log N)$

Dado un nodo $u$, sabemos que cualquier camino en el árbol o bien pasa por $u$ o bien está completamente contenido en uno de los subárboles de los hijos de $u$. Esto significa que para resolver el problema en un árbol que contiene $u$, podemos:

1. Contar el número de caminos de longitud $k$ que pasan por $u$.
2. Quitar $u$ y sus aristas incidentes del árbol.
3. Resolver el problema de forma recursiva en los nuevos árboles que se forman.

Contar caminos que pasan por $u$

Podemos contar el número de caminos que pasan por $u$ en un árbol de tamaño $M$ en tiempo $\mathcal O(M)$.

Procesamos cada uno de los subárboles de los hijos de $u$ en orden. Sea $\texttt{cnt}_i[d]$ el número de nodos que hemos visitado con profundidad $d$ después de procesar los primeros $i$ subárboles de los hijos de $u$. Un nodo con profundidad $d$ en el $i$-ésimo subárbol de un hijo de $u$ contribuirá $\texttt{cnt}_{i - 1}[k - d]$ caminos a la respuesta.

Usando dos DFS para cada hijo, podemos calcular su contribución y luego actualizar $\texttt{cnt}$ en tiempo total $\mathcal O(M)$.

Aplicar descomposición por centroide

Podemos combinar la idea anterior con descomposición por centroide para resolver todo el problema en tiempo $\mathcal O(N \log N)$.

Al procesar algún árbol, simplemente hacemos que el $u$ que usamos sea el centroide de ese árbol. Como el árbol de centroides tiene profundidad $\mathcal O(\log N)$ y cada capa del árbol de centroides toma tiempo amortizado $\mathcal O(N)$ para procesarse, la complejidad total es $\mathcal O(N \log N)$.

Implementación

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
 
int n, k;
vector<int> graph[200001];
int subtree[200001];
 
ll ans = 0;
int cnt[200001]{1}, mx_depth;
bool processed[200001];
 
int get_subtree_sizes(int node, int parent = 0) {
	subtree[node] = 1;
	for (int i : graph[node])
		if (!processed[i] && i != parent) subtree[node] += get_subtree_sizes(i, node);
	return subtree[node];
}
 
int get_centroid(int desired, int node, int parent = 0) {
	for (int i : graph[node])
		if (!processed[i] && i != parent && subtree[i] >= desired)
			return get_centroid(desired, i, node);
	return node;
}
 
void get_cnt(int node, int parent, bool filling, int depth = 1) {
	if (depth > k) return;
	mx_depth = max(mx_depth, depth);
	if (filling) cnt[depth]++;
	else ans += cnt[k - depth];
	for (int i : graph[node])
		if (!processed[i] && i != parent) get_cnt(i, node, filling, depth + 1);
}
 
void centroid_decomp(int node = 1) {
	int centroid = get_centroid(get_subtree_sizes(node) >> 1, node);
	processed[centroid] = true;
	mx_depth = 0;
	for (int i : graph[centroid])
		if (!processed[i]) {
			get_cnt(i, centroid, false);
			get_cnt(i, centroid, true);
		}
	fill(cnt + 1, cnt + mx_depth + 1, 0);
	for (int i : graph[centroid])
		if (!processed[i]) centroid_decomp(i);
}
 
int main() {
	cin.tie(0)->sync_with_stdio(0);
	cin >> n >> k;
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		graph[u].push_back(v);
		graph[v].push_back(u);
	}
	centroid_decomp();
	cout << ans;
	return 0;
}