Cod sursa(job #1770734)

Utilizator Andrei1998Andrei Constantinescu Andrei1998 Data 4 octombrie 2016 19:36:57
Problema Cerere Scor 100
Compilator cpp Status done
Runda Arhiva de probleme Marime 2.04 kb
#include <fstream>
#include <vector>

using namespace std;

class InputReader {
    public:
        InputReader() {}
        InputReader(const char *file_name) {
            input_file = fopen(file_name, "r");
            cursor = 0;
            fread(buffer, SIZE, 1, input_file);
        }
        inline InputReader &operator >>(int &n) {
            while(buffer[cursor] < '0' || buffer[cursor] > '9') {
                advance();
            }
            n = 0;
            while('0' <= buffer[cursor] && buffer[cursor] <= '9') {
                n = n * 10 + buffer[cursor] - '0';
                advance();
            }
            return *this;
        }
    private:
        FILE *input_file;
        static const int SIZE = 1 << 17;
        int cursor;
        char buffer[SIZE];
        inline void advance() {
            ++ cursor;
            if(cursor == SIZE) {
                cursor = 0;
                fread(buffer, SIZE, 1, input_file);
            }
        }
};

const int NMAX = 100000 + 5;

vector <int> tree[NMAX];
vector <int> tree2[NMAX];

int k[NMAX];
int ans[NMAX];

int stk[NMAX];
int stkSz;

int father[NMAX];

void dfs(int node) {
    if (k[node])
        tree2[stk[stkSz - k[node]]].push_back(node);
    else
        tree2[0].push_back(node);

    for (auto it: tree[node]) {
        stk[++ stkSz] = it;
        dfs(it);
        -- stkSz;
    }
}

int h[NMAX];
void dfs2(int node) {
    for (auto it: tree2[node]) {
        h[it] = 1 + h[node];
        dfs2(it);
    }
}

int main()
{
    InputReader cin("cerere.in");
    ofstream cout("cerere.out");

    int n;
    cin >> n;

    for (int i = 1; i <= n; ++ i)
        cin >> k[i];


    for (int i = 1; i < n; ++ i) {
        int a, b;
        cin >> a >> b;
        father[b] = a;
        tree[a].push_back(b);
    }

    int one = 1;
    for (int i = 1; i <= n; ++ i)
        if (!father[i])
            one = i;

    stk[++ stkSz] = one;
    dfs(one);
    dfs2(0);

    for (int i = 1; i <= n; ++ i)
        cout << h[i] - 1 << " \n"[i == n];

    //cin.close();
    cout.close();
    return 0;
}