Select Edges

Explicación

Para resolver este problema, escribimos una DP en árboles. El único factor que necesitamos para diferenciar nodos es si podemos o no conectar un nodo dado con su padre. Así, nuestro estado de DP es el siguiente:

$\texttt{dp}[u][0]$ es el mejor resultado en el subárbol de $u$ si podemos conectar $u$ con su padre
$\texttt{dp}[u][1]$ es el mejor resultado en el subárbol de $u$ si no podemos conectar $u$ con su padre

Para un nodo dado $u$ , la “ganancia” que obtenemos al agregar este nodo es:

(w + \texttt{dp}[u][0]) - \texttt{dp}[u][1]

Así, queremos usar nuestras $d[u]$ aristas permitidas en los nodos que más beneficio dan al agregarles una arista. Notamos que manejamos los casos de $\texttt{dp}[u][0])$ y $\texttt{dp}[u][1]$ de forma bastante similar.

Implementación

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

C++


#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
 
int main() {
	int n;
	cin >> n;
 
	vector<int> d(n);
	for (int &i : d) { cin >> i; }
 
	vector<vector<array<int, 2>>> adj(n);
	for (int i = 0; i < n - 1; i++) {
		int u, v, w;
		cin >> u >> v >> w;
		u--, v--;
		adj[u].push_back({v, w});
		adj[v].push_back({u, w});
	}
 
	// El estado de DP es para la relación con su nodo padre
	// {puede conectar, no puede conectar}
	vector<array<ll, 2>> dp(n);
	const auto dfs = [&](int u, int p, auto self) -> void {
		auto &[can_con, no_con] = dp[u];
		vector<ll> diffs;  // diferencia entre usar la arista y no usarla
 
		for (const auto &[v, w] : adj[u]) {
			if (v == p) { continue; }
			self(v, u, self);
			can_con += dp[v][1];
			no_con += dp[v][1];
			if (d[v] > 0) { diffs.push_back(w + dp[v][0] - dp[v][1]); }
		}
 
		int leftover = d[u];
 
		sort(begin(diffs), end(diffs), greater());
		for (const ll &val : diffs) {
			if (leftover == 0 || val < 0) { break; }
			if (leftover > 1) {
				can_con += val;
				no_con += val;
			} else {
				no_con += val;
			}
			leftover--;
		}
	};
	dfs(0, -1, dfs);
 
	cout << dp[0][1] << endl;
}