Skip to Content
CTRL K

Tree Distances II

Complejidad temporal: O(n)\mathcal O(n)

Es fácil hallar la suma de distancias desde un solo nodo: basta enraizar el árbol en ese nodo, hacer un DFS y sumar las profundidades de cada otro nodo a la respuesta. Lamentablemente, nn puede llegar hasta 2⋅1052 \cdot 10^5, así que no podemos hacer esto para cada nodo.

Si tenemos la respuesta para algún nodo (digamos el nodo 1), ¿cómo podemos hallar rápidamente la respuesta para sus vecinos?

La observación clave es que si reenraizamos (rerooting) el árbol en el vecino del nodo 1 (digamos el nodo 2), entonces

  • Las profundidades de todos los nodos del subárbol del nodo 2 disminuyen en 1.
  • Las profundidades de todos los nodos fuera de su subárbol aumentan en 1.

Esto nos da una forma elegante de pasar de la respuesta del nodo 1 a la del nodo 2 usando solo nn y el tamaño del subárbol del nodo 2. Observemos que el cambio en la respuesta es exactamente n−2(node 2’s subtree size)n - 2(\text{node 2's subtree size}).

Podemos usar DFS para hallar tanto la respuesta del nodo 1 como el tamaño del subárbol de cada nodo cuando el árbol está enraizado en el nodo 1, y luego otro DFS para computar todas las respuestas.

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
 
int n;
vector<int> graph[200001];
ll dp[200001], ans[200001];
 
void dfs1(int node = 1, int parent = 0, ll depth = 0) {
	ans[1] += depth;
	dp[node] = 1;
	for (int i : graph[node])
		if (i != parent) {
			dfs1(i, node, depth + 1);
			dp[node] += dp[i];
		}
}
 
void dfs2(int node = 1, int parent = 0) {
	for (int i : graph[node])
		if (i != parent) {
			ans[i] = ans[node] + n - 2 * dp[i];
			dfs2(i, node);
		}
}
 
int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cin >> n;
	for (int i = 1; i < n; i++) {
		int a, b;
		cin >> a >> b;
		graph[a].push_back(b);
		graph[b].push_back(a);
	}
	dfs1();
	dfs2();
	for (int i = 1; i <= n; i++) cout << ans[i] << ' ';
	return 0;
}
Fundamentos