Cod sursa(job #2922897)

Utilizator tiut_cristianTiut Cristian tiut_cristian Data 10 septembrie 2022 16:03:20
Problema Asmax Scor 60
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.65 kb
#include <fstream>
#include <vector>
#include <stack>
#include <algorithm>

using namespace std;

ifstream fin("asmax.in");
ofstream fout("asmax.out");

void citire(int & n, vector<int> & v, vector<vector<int>> & adj)
{
    fin >> n;
    v.resize(n+1);
    for(int i = 1; i <= n; i++)
        fin >> v[i];
    adj.resize(n+1);
    for(int i = 1; i <= n-1; i++)
    {
        int a, b;
        fin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
}

int dfs_it(int n, const vector<int> & v, const vector<vector<int>> & adj)
{
    vector<int> nivel;
    nivel.resize(n+1, -1);
    nivel[1] = 0;

    stack<pair<int, int>> st;
    st.push({1, 0});

    vector<int> d(n+1);

    while(!st.empty())
    {
        int node = st.top().first;
        int index = st.top().second;
        if(index < adj[node].size())
        {
            int copil = adj[node][index];
            if(nivel[copil] == -1)
            {
                nivel[copil] = nivel[node] + 1;
                st.push({copil, 0});
            }
            st.top().second++;
        }
        else
        {
            st.pop();
            d[node] = v[node];
            for(auto j : adj[node])
                if(nivel[j] == nivel[node] + 1 && d[j] > 0)
                    d[node] += d[j];
        }
    }
/*
    int maxim = d[1];
    for(int i = 2; i <= n; i++)
        if(d[i] > maxim)
            maxim = d[i];*/
    return *max_element(d.begin() + 1, d.end());
}

int main()
{
    int n;
    vector<int> v;
    vector<vector<int>> adj;

    citire(n, v, adj);
    fout << dfs_it(n, v, adj)<< '\n';
    return 0;
}