Problem p02728 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;
#define ll long long
#define ld long double
#define f first
#define s second

const int N = 2e5 + 5;
int n, fact[N], inv[N], dp[N], ans[N], sz[N], mod = 1e9 + 7;
vector<int> adj[N], preDp[N], sufDp[N], preInv[N], sufInv[N];

int mul(int a, int b) {
  a = (a + mod) % mod;
  b = (b + mod) % mod;
  return (1ll * a * b) % mod;
}

int add(int a, int b) {
  a = (a + mod) % mod;
  b = (b + mod) % mod;
  return (a + b) % mod;
}

int fastpow(int base, int power) {
  if (!power)
    return 1;
  int ret = fastpow(base, power / 2);
  ret = mul(ret, ret);
  if (power & 1)
    ret = mul(ret, base);
  return ret;
}

void pree() {
  fact[0] = 1, inv[0] = 1;
  for (int i = 1; i < N; i++) {
    fact[i] = mul(i, fact[i - 1]);
    inv[i] = fastpow(fact[i], mod - 2);
  }
}

void dfs(int node, int par) {
  sz[node] = 1, dp[node] = 1;
  for (auto ch : adj[node]) {
    if (ch != par) {
      dfs(ch, node);
      sz[node] += sz[ch];
      dp[node] = mul(dp[node], dp[ch]);
    }
  }
  dp[node] = mul(dp[node], fact[sz[node] - 1]);
  for (auto ch : adj[node]) {
    if (ch != par)
      dp[node] = mul(dp[node], inv[sz[ch]]);
  }
}

void ylla(int node) {
  preDp[node] = vector<int>(adj[node].size() + 2, 0);
  sufDp[node] = vector<int>(adj[node].size() + 2, 0);
  preInv[node] = vector<int>(adj[node].size() + 2, 0);
  sufInv[node] = vector<int>(adj[node].size() + 2, 0);

  for (int i = 0; i < adj[node].size(); i++) {
    int ch = adj[node][i];
    preDp[node][i] = dp[ch];
    preInv[node][i] = inv[sz[ch]];
    if (i)
      preDp[node][i] = mul(preDp[node][i], preDp[node][i - 1]),
      preInv[node][i] = mul(preInv[node][i], preInv[node][i - 1]);
  }

  for (int i = adj[node].size() - 1; i >= 0; i--) {
    int ch = adj[node][i];
    sufDp[node][i] = dp[ch];
    sufInv[node][i] = inv[sz[ch]];
    if (i < adj[node].size() - 1)
      sufDp[node][i] = mul(sufDp[node][i], sufDp[node][i + 1]),
      sufInv[node][i] = mul(sufInv[node][i], sufInv[node][i + 1]);
  }
}

void solve(int node, int par, int d, int in) {
  int tmpPar = dp[par], tmpNode = dp[node];

  sz[par] -= sz[node];
  dp[par] = fact[sz[par] - 1];
  dp[par] = mul(dp[par], d),
  dp[par] = mul(dp[par], in); // el for loop elle shelnaha

  sz[node] += sz[par];
  dp[node] = mul(1, fact[sz[node] - 1]);
  for (auto ch : adj[node]) {
    dp[node] = mul(dp[node], dp[ch]);
    dp[node] = mul(dp[node], inv[sz[ch]]);
  }
  ans[node] = dp[node];

  ylla(node);

  for (int i = 0; i < adj[node].size(); i++) {
    int ch = adj[node][i], d = 1, in = 1;
    if (ch == par)
      continue;
    if (i)
      d = mul(d, preDp[node][i - 1]), in = mul(in, preInv[node][i - 1]);
    if (i < adj[node].size() - 1)
      d = mul(d, sufDp[node][i + 1]), in = mul(in, sufInv[node][i + 1]);
    solve(ch, node, d, in);
  }

  sz[node] -= sz[par];
  sz[par] += sz[node];
  dp[node] = tmpNode, dp[par] = tmpPar;
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  pree();
  cin >> n;
  for (int i = 1; i < n; i++) {
    int x, y;
    cin >> x >> y;
    adj[x].push_back(y);
    adj[y].push_back(x);
  }

  dfs(1, 1);
  ans[1] = dp[1];

  ylla(1);
  for (int i = 0; i < adj[1].size(); i++) {
    int ch = adj[1][i], d = 1, in = 1;
    if (i)
      d = mul(d, preDp[1][i - 1]), in = mul(in, preInv[1][i - 1]);
    if (i < adj[1].size() - 1)
      d = mul(d, sufDp[1][i + 1]), in = mul(in, sufInv[1][i + 1]);
    solve(ch, 1, d, in);
  }

  for (int i = 1; i <= n; i++)
    cout << ans[i] << '\\n';

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
#define ll long long
#define ld long double
#define f first
#define s second

const int N = 2e5 + 5;
int n, fact[N], inv[N], dp[N], ans[N], sz[N], mod = 1e9 + 7;
vector<int> adj[N];
bool vis[N];

int fastpow(int base, int power) {
  if (!power)
    return 1;
  int ret = fastpow(base, power / 2);
  ret = (1ll * ret * ret) % mod;

  if (power & 1)
    ret = (1ll * ret * base) % mod;

  return ret;
}

void pre(int n) {
  fact[0] = 1, inv[0] = 1;
  for (int i = 1; i <= n; i++) {
    fact[i] = (1ll * i * fact[i - 1]) % mod;
    inv[i] = fastpow(fact[i], mod - 2);
  }
}

void dfs(int node, int par) {
  if (vis[node])
    return;
  vis[node] = 1;

  sz[node] = 1, dp[node] = 1;
  for (auto ch : adj[node]) {
    if (ch != par) {
      dfs(ch, node);
      sz[node] += sz[ch];
      dp[node] = (1ll * dp[node] * dp[ch]) % mod;
    }
  }
  dp[node] = (1ll * dp[node] * fact[sz[node] - 1]) % mod;

  for (auto ch : adj[node]) {
    if (ch != par)
      dp[node] = (1ll * dp[node] * inv[sz[ch]]) % mod;
  }
}

void solve(int node, int par) {
  if (vis[node])
    return;
  vis[node] = 1;

  int tmpPar = dp[par], tmpNode = dp[node];

  sz[par] -= sz[node];
  sz[node] += sz[par];
  dp[par] = fact[sz[par] - 1];

  for (auto ch : adj[par]) {
    if (ch != node) {
      dp[par] = (1ll * dp[par] * dp[ch]) % mod;
      dp[par] = (1ll * dp[par] * inv[sz[ch]]) % mod;
    }
  }

  dp[node] = fact[sz[node] - 1];
  for (auto ch : adj[node]) {
    dp[node] = (1ll * dp[node] * dp[ch]) % mod;
    dp[node] = (1ll * dp[node] * inv[sz[ch]]) % mod;
    if (ch != par)
      solve(ch, node);
  }
  ans[node] = dp[node];

  sz[node] -= sz[par];
  sz[par] += sz[node];
  dp[node] = tmpNode, dp[par] = tmpPar;
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  scanf("%d", &n);
  int x, y;
  pre(n);
  for (int i = 1; i < n; i++) {
    scanf("%d%d", &x, &y);
    adj[x].push_back(y);
    adj[y].push_back(x);
  }

  dfs(1, 1);
  ans[1] = dp[1];
  memset(vis, 0, sizeof(vis));
  vis[1] = 1;

  for (auto ch : adj[1])
    solve(ch, 1);

  for (int i = 1; i <= n; i++)
    printf("%d\\n", ans[i]);

  return 0;
}