Color the Tree

Explicación

Solución naive

Observamos que cualquier operación solo afectará nodos de la misma profundidad. Así, en lugar de resolver para todas las profundidades a la vez, podemos tomar la suma del costo mínimo para rellenar cada profundidad.

Consideremos el conjunto de nodos a profundidad $d$ relativa a la raíz. Definamos $\texttt{dp}[v]$ como el costo mínimo para rellenar todos los vértices de profundidad $d$ en el subárbol de $v$ , considerando solo operaciones que seleccionan vértices en $v$ o más abajo.

Si $v$ es de profundidad $d$ (relativa a la raíz), entonces debemos rellenarlo con costo $a_0$

\texttt{dp}[v] =a_0

En caso contrario, podemos o bien incurrir un costo de $a_{d-\text{depth}[v]}$ operando sobre $v$ , o no operar sobre $v$ y en cambio resolver para cada subárbol:

\texttt{dp}[v] = \min\left( a_{d-\text{depth}[v]},\sum_{u\in\text{Children}(v)} \text{dp}[u]\right)

Nuestra respuesta será la suma de $\texttt{dp}[0]$ para cada valor de $d$ , lo que da una solución en tiempo $\mathcal{O}(n^2)$ .

Optimización con árbol virtual

Una observación que podemos hacer es que, para una cadena larga que termina en $v$ , podemos procesarla rápido tomando

\texttt{dp}[v] = \min \left( \min_{x\in[d-\text{depth}[v],d]}a_x, \sum_{u\in\text{Children}(v)} \text{dp}[u] \right)

Apuntaremos a aprovechar este hecho. Definamos $S_d$ como el conjunto de vértices a profundidad $d$ , y tomemos $S$ como el árbol virtual de $S_d$ .

Continuamos de forma similar a antes, pero solo calculamos valores de programación dinámica para nodos en $S$ . La consulta de mínimo en un rango se puede responder usando estructuras de datos en tiempo logarítmico o lineal.

\texttt{dp}[v] = \min \left( \min_{x\in[d-\text{depth}[v],d]}a_x, \sum_{u\in\text{Virtual Children}(v)} \text{dp}[u] \right)

Implementación

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


#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
 
constexpr int MAX_N = 1e5 + 1;
constexpr int LG = 17;
 
int tin[MAX_N], tout[MAX_N], d[MAX_N], lift[MAX_N][LG], timer;
ll dp[MAX_N];
vector<int> adj[MAX_N], vadj[MAX_N], at_d[MAX_N];
 
void dfs(int v, int p) {
	d[v] = v ? d[p] + 1 : 0;
	at_d[d[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(vector<int> &key) {
	// construct virtual tree from set of nodes
	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(key[i - 1], key[i]));
	}
 
	sort(res.begin(), res.end(), sort_tin);
	res.erase(unique(res.begin(), res.end()), res.end());
 
	for (int v : res) { vadj[v].clear(); }
 
	vector<int> stk = {res[0]};
 
	for (int i = 1; i < (int)res.size(); i++) {
		while (tin[res[i]] > tout[stk.back()]) { stk.pop_back(); }
 
		vadj[stk.back()].push_back(res[i]);
		stk.push_back(res[i]);
	}
 
	return res;
}
 
// BeginCodeSnip{RMQ Template}
template <class T> class MinSegmentTree {
  private:
	const T DEFAULT = numeric_limits<T>().max();
 
	vector<T> segtree;
	int len;
 
  public:
	MinSegmentTree(int len) : len(len), segtree(len * 2, DEFAULT) {}
 
	void set(int ind, T val) {
		ind += len;
		segtree[ind] = val;
		for (; ind > 1; ind /= 2) {
			segtree[ind / 2] = min(segtree[ind], segtree[ind ^ 1]);
		}
	}
 
	T range_min(int start, int end) {
		T res = DEFAULT;
		for (start += len, end += len; start < end; start /= 2, end /= 2) {
			if (start % 2 == 1) { res = min(res, segtree[start++]); }
			if (end % 2 == 1) { res = min(res, segtree[--end]); }
		}
		return res;
	}
};
// EndCodeSnip
 
int main() {
	int test_num;
	cin >> test_num;
	for (int t = 0; t < test_num; t++) {
		int n;
		cin >> n;
 
		for (int i = 0; i <= n; i++) {
			adj[i].clear();
			at_d[i].clear();
		}
 
		MinSegmentTree<int> rmq(n);
		for (int i = 0; i < n; i++) {
			int a;
			cin >> a;
			rmq.set(i, a);
		}
 
		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);
 
		ll ans = 0;
 
		for (int dep = 0; dep < n; dep++) {
			vector<int> vt = vtree(at_d[dep]);
			reverse(vt.begin(), vt.end());
 
			for (int v : vt) {
				dp[v] = rmq.range_min(dep - d[v], dep + 1);
 
				if (!vadj[v].empty()) {
					ll sum = 0;
					for (int u : vadj[v]) { sum += dp[u]; }
					dp[v] = min(dp[v], sum);
				}
			}
 
			if (!vt.empty()) { ans += dp[vt.back()]; }
		}
 
		cout << ans << "\n";
	}
}