Cod sursa(job #3272171)

Utilizator iustincmbMaican Iustin iustincmb Data 28 ianuarie 2025 18:03:50
Problema Cerere Scor 0
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.16 kb
#include <bits/stdc++.h>
using namespace std;
vector<int> adj[100005];
int v[100005];
int st[100005], l;
int f[100005];
bool vis[100005];
void dfs(int node){
    vis[node]=1;
    st[++l]=node;
    f[node]=st[l-v[node]];
    for(auto i : adj[node]){
        if(vis[i]==0)
            dfs(i);
    }
    st[--l];
}
int ras[100005];
int main(){
    ifstream cin ("cerere.in");
    ofstream cout ("cerere.out");
    int n;
    cin >> n;
    for(int i=1; i<=n; i++)
        cin >> v[i];
    int a, b;
    for(int i=1; i<=n-1; i++){
        cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    dfs(1);
    for(int i=1; i<=n; i++)
        ras[i]=-1;
    for(int i=1; i<=n; i++){
        if(v[i]==0)
            ras[i]=0;
    }
    for(int i=1; i<=n; i++){
        if(v[i]!=0){
            int poz=i, cnt=0;
            while(ras[poz]==-1){
                poz=f[poz];
                cnt++;
            }
            cnt+=ras[poz];
            poz=i;
            while(ras[poz]==-1){
                ras[poz]=cnt;
                poz=f[poz];
                cnt--;
            }
        }
    }
    for(int i=1; i<=n; i++)
        cout << ras[i] << ' ';
    return 0;
}