Tree Painting

Editorial oficial (C++) 

Explicación

Nótese que elegir un nodo hoja es más óptimo que elegir algo en el medio, porque así tenemos el tamaño entero del resto del árbol como parte del componente conexo restante. Esto sugiere una solución de DP en árboles, donde elegimos un nodo hoja, subimos a través de los ancestros de esa hoja y bajamos por cada uno de los hermanos de nuestros ancestros, incorporando el costo máximo posible para cada uno de sus subárboles. Por ejemplo, para el siguiente árbol:

Si seleccionamos el 4.º nodo como punto de partida, entonces el proceso se vería así:

1. Subimos hasta 3
2. Bajamos hasta 5
3. Subimos hasta 1
4. Bajamos hasta 2

Hay que mantener el costo de dos escenarios distintos para cada nodo: trabajar hacia arriba desde un nodo hoja hasta el nodo actual y trabajar hacia abajo desde el nodo actual a lo largo de todo su subárbol, lo que significa que hay que mantener 2 arreglos de DP. Definamos las transiciones de DP para cada uno de los arreglos. Usaremos $dp_1$ para el costo de trabajar hacia abajo y $dp_2$ para el costo de trabajar hacia arriba, y $sz$ para el tamaño de un subárbol.

Podemos calcular este arreglo de DP con una sola búsqueda DFS. Ahora, para el segundo arreglo de DP:

Aún hay un problema. El costo de $sz[u] - sz[c]$ se multiplica por el número de nodos entre el nodo actual y el nodo inicial, inclusive. Sin embargo, no podemos anticipar este número de antemano, y podría haber algún nodo ancestro más adelante que habría preferido una distancia mayor sobre un costo mayor. Esto significa que hay que incorporar el $sz[u] - sz[c]$ de cada nodo ancestro al inicializar el valor de $dp_2$ para los nodos hoja. De esa forma, se incorpora al costo de antemano y no hay que preocuparse por elegir el nodo inicial incorrecto con un costo final no óptimo.

Podemos calcular estos valores iniciales con otra búsqueda DFS antes de la que calcula los valores de $dp_2$. Así podemos quitar el término que incluye $sz[u] - sz[c]$ de la fórmula de recurrencia de $dp_2$. Para calcular estos valores iniciales, notamos que cada vez que bajamos un nivel de profundidad, sumamos el tamaño de todos los hermanos y el padre combinados dos veces, una por el padre y una por el nodo actual. También sumamos el costo de todos los hermanos de los ancestros y el padre combinados una vez más. Digamos que $p_2[u]$ es el costo inicial para un nodo hoja $u$, y $p_1[u]$ es el costo de los hermanos de los ancestros y el padre combinados hasta $u$. Entonces podemos calcular $p_1[u]$ y $p_2[u]$ usando la siguiente recurrencia:

Implementación

Complejidad temporal: $\mathcal{O}(N)$

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
 
vector<vector<int>> g;
vector<ll> dp1, dp2;
vector<ll> p1, p2;
vector<ll> sz;
 
// Calculate sz and dp1
void dfs1(int u, int p) {
	sz[u] = 1;
	for (int &v : g[u]) {
		if (v == p) { continue; }
		dfs1(v, u);
		sz[u] += sz[v];
		dp1[u] += dp1[v];
	}
	dp1[u] += sz[u];
}
 
// Calculate cost incorporated for initial value of dp2
void dfs2(int u, int p) {
	if (p != -1) {
		p1[u] = p1[p] + sz[p] - sz[u];
		p2[u] = p2[p] + p1[u] + sz[p] - sz[u];
	}
	for (int &v : g[u]) {
		if (v == p) { continue; }
		dfs2(v, u);
	}
}
 
// Calculate dp2
void dfs3(int u, int p) {
	dp2[u] = (g[u].size() == 1 ? p2[u] + 1 : 0);
	ll tot = 0;
	for (int &v : g[u]) {
		if (v == p) { continue; }
		tot += dp1[v];
	}
	for (int &v : g[u]) {
		if (v == p) { continue; }
		dfs3(v, u);
		if (dp2[v] + tot - dp1[v] > dp2[u]) { dp2[u] = dp2[v] + tot - dp1[v]; }
	}
}
 
int main() {
	int n;
	cin >> n;
 
	g.resize(n);
	dp1.resize(n);
	dp2.resize(n);
	sz.resize(n);
	p1.resize(n);
	p2.resize(n);
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y;
		g[--x].push_back(--y);
		g[y].push_back(x);
	}
 
	// Run first DFS to calculate sz and dp1
	dfs1(0, -1);
 
	// Run second DFS to calculate p1 and p2
	dfs2(0, -1);
 
	// Run third DFS to calculate dp2
	dfs3(0, -1);
 
	// Find maximum dp1/dp2 among all nodes
	ll ans = 0;
	for (int i = 0; i < n; ++i) { ans = max(ans, max(dp1[i], dp2[i])); }
 
	cout << ans << endl;
}