543D - Road Improvement - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;

const int N = 2e5 + 100, MOD = 1e9 + 7;

int cnt[N], ans[N], l[N], r[N], res[N];
vector<int> adj[N];
int n;

int mul(int a, int b) {
    return 1LL * a * b % MOD;
}

int sum(int a, int b) {
    a += b;
    if (a >= MOD)
	a -= MOD;
    return a;
}

void dfs(int v) {
    cnt[v] = 1;
    for (auto u : adj[v]) {
	dfs(u);
	cnt[v] = mul(cnt[v], cnt[u] + 1);
    }
}

void DFS(int v, int num) {
    ans[v] = cnt[v];
    int sz = adj[v].size();
    l[0] = r[sz] = 1;
    for (int i = 1, j = sz - 1; i <= sz; i++, j--) {
	l[i] = mul(cnt[adj[v][i - 1]] + 1, l[i - 1]);
	r[j] = mul(cnt[adj[v][j]] + 1, r[j + 1]);
    }
    for (int i = 0; i < sz; i++)
	res[adj[v][i]] = mul(l[i], r[i + 1]);
    for (int i = 0; i < sz; i++) {
	int u = adj[v][i];
	int tmp = mul(num, res[u]) + 1;
	cnt[u] = mul(cnt[u], tmp);
	DFS(u, tmp);
    }
}

int main() {
    ios_base::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    cin >> n;
    for (int i = 1; i < n; i++) {
	int p;
	cin >> p;
	adj[--p].push_back(i);
    }
    dfs(0);
    DFS(0, 1);
    for (int i = 0; i < n; i++)
	cout << ans[i] << " ";
}