Cod sursa(job #2709789)

Utilizator mex7Alexandru Valentin mex7 Data 21 februarie 2021 11:55:53
Problema Cerere Scor 100
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.38 kb
#include <bits/stdc++.h>
#define ll long long
#define cin fin
#define cout fout
using namespace std;

ifstream fin("cerere.in");
ofstream fout("cerere.out");
vector <int> children[100005];
int n, father[100005], want[100005], wantedAncestor[100005], grade[100005], numOfMonkeys[100005];
bitset <100005> reached;

int getRoot(int monkey) {
    if (father[monkey] == 0)
        return monkey;
    return getRoot(father[monkey]);
}

void findAncestors(int monkey, int length) {
    grade[length] = monkey;
    wantedAncestor[monkey] = grade[length - want[monkey]];

    for (int child : children[monkey])
        findAncestors(child, length + 1);
}

int findNumOfMonkeys(int monkey) {
    if (reached[monkey])
        return numOfMonkeys[monkey];

    reached[monkey] = 1;
    if (wantedAncestor[monkey] == monkey)
        return numOfMonkeys[monkey] = 0;
    return numOfMonkeys[monkey] = 1 + findNumOfMonkeys(wantedAncestor[monkey]);
}

int main() {
    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> want[i];
    for (int i = 1; i < n; i++) {
        int monkey1, monkey2;
        cin >> monkey1 >> monkey2;
        father[monkey2] = monkey1;
        children[monkey1].push_back(monkey2);
    }

    findAncestors(getRoot(1), 1);

    for (int monkey = 1; monkey <= n; monkey++)
        cout << findNumOfMonkeys(monkey) << ' ';

    return 0;
}