Cod sursa(job #2922357)

Utilizator BeilandArnoldArnold Beiland BeilandArnold Data 8 septembrie 2022 01:39:31
Problema Asmax Scor 100
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.71 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 un subarbore 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<int> level(n+1, -1);
    
    st.push({ 1, 0 });
    level[1] = 0;

    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) {
            dp[node] = v[node];

            for (int i = 0; i < nadj; ++i) {
                int child = adj[node][i];
                if (level[child] == level[node] + 1 && dp[child] > 0)
                    dp[node] += dp[child];
            }

            st.pop();
        }
        else {
            int child = adj[node][index];
            ++index;

            if (level[child] == -1) {
                st.push({ child, 0 });
                level[child] = level[node] + 1;
            }
        }
    }

    out << *max_element(dp.begin()+1, dp.end()) << '\n';

    return 0;
}