Cod sursa(job #2922055)

Utilizator BeilandArnoldArnold Beiland BeilandArnold Data 3 septembrie 2022 14:07:18
Problema Asmax Scor 10
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 2.06 kb
// 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;
}