Pagini recente » Cod sursa (job #1212219) | Cod sursa (job #2442875) | Cod sursa (job #1932067) | Cod sursa (job #748899) | Cod sursa (job #2922055)
// idea: dynamic programming on a tree
#include <fstream>
#include <vector>
#include <stack>
#include <algorithm>
using namespace std;
int main()
{
ifstream in("asmax.in");
ofstream out("asmax.out");
int n;
in >> n;
vector<int> v(n+1, 0);
for (int i = 1; i <= n; ++i)
in >> v[i];
vector<vector<int>> adj(n+1);
for (int i = 0; i < n-1; ++i) {
int x, y;
in >> x >> y;
adj[x].push_back(y);
adj[y].push_back(x);
}
vector<int> dp(n+1); // dp[i] - valoarea maximă pentru subarborele cu
// rădăcina în i
stack<pair<int,int>> st; // .first - nodul
// .second - următorul index care poate fi folosit
// din lista lui de adiacență
vector<bool> vizitat(n+1, false);
st.push({ 1, 0 });
vizitat[1] = true;
while (!st.empty()) {
pair<int, int> &pair = st.top();
int &node = pair.first;
int &index = pair.second;
int nadj = adj[node].size();
if (index == nadj) {
if (nadj == 0)
dp[node] = v[node];
else {
// fie luăm doar copilul cu valoarea maximă
// fie suma copiilor pozitivi + valoarea nodului curent
int copil_max = dp[adj[node][0]];
int smax = v[node] + max(0, copil_max);
for (int i = 1; i < nadj; ++i) {
int val = dp[adj[node][i]];
copil_max = max(copil_max, val);
if (val > 0)
smax += val;
}
dp[node] = max(copil_max, smax);
}
st.pop();
}
else {
int copil = adj[node][index];
++index;
if (!vizitat[copil]) {
st.push({ copil, 0 });
vizitat[copil] = true;
}
}
}
out << dp[1] << '\n';
return 0;
}