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

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
vector<int> cntd;
vector<int> sz;
vector<int> cntu;
vector<int> par;
vector<vector<int>> adj;
vector<int> fact;
vector<int> inv;
vector<int> ans;
const int MOD = 1e9 + 7;

void mul(int &a, int b) { a = 1ll * a * b % MOD; }

void add(int &a, int b) {
  a += b;
  if (a >= MOD)
    a -= MOD;
}

ll pow1(ll n, ll p) {
  ll r = 1, c = n;
  while (p) {
    if (p & 1)
      (r *= c) %= MOD;
    p /= 2;
    (c *= c) %= MOD;
  }
  return r;
}

void dfs_d(int u, int p = -1) {
  sz[u] = 1;
  cntd[u] = 1;
  par[u] = p;
  for (auto el : adj[u]) {
    if (el == p)
      continue;
    dfs_d(el, u);
    sz[u] += sz[el];
    mul(cntd[u], cntd[el]);
    mul(cntd[u], inv[sz[el]]);
  }
  mul(cntd[u], fact[sz[u] - 1]);
}
int n;

void dfs_u(int u, int p = -1) {
  if (p == -1)
    cntu[u] = 1;
  else {
    cntu[u] = cntd[p];
    int szp = n - sz[u];
    mul(cntu[u], inv[sz[p] - 1]);
    mul(cntu[u], fact[n - sz[u] - 1]);
    mul(cntu[u], fact[sz[u]]);
    mul(cntu[u], pow1(cntd[u], MOD - 2));
    int szpp = n - sz[p];
    mul(cntu[u], cntu[p]);
    mul(cntu[u], inv[szpp]);
  }
  for (auto el : adj[u]) {
    if (el == p)
      continue;
    dfs_u(el, u);
  }
}

int main() {
  //  freopen("in.txt", "r", stdin);
  //  freopen("out.txt", "w", stdout);
  cin >> n;
  cntd.resize(n);
  sz.resize(n);
  cntu.resize(n);
  par.resize(n);
  adj.resize(n);
  fact.resize(n + 1);
  inv.resize(n + 1);
  ans.resize(n);
  fact[0] = inv[0] = 1;
  for (int i = 1; i <= n; i++) {
    fact[i] = 1ll * fact[i - 1] * i % MOD;
    inv[i] = pow1(fact[i], MOD - 2);
  }
  //  assert(1ll * fact[3] * inv[3] % MOD == 1);

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    scanf("%d %d", &u, &v);
    u--, v--;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  dfs_d(0);
  dfs_u(0);
  //  for(int i = 0;i < n;i++)
  //    cout << cntu[i] << endl;
  for (int i = 0; i < n; i++) {
    int szp = n - sz[i];
    int res = cntu[i];
    mul(res, inv[szp]);
    for (auto el : adj[i]) {
      if (el == par[i])
        continue;
      mul(res, cntd[el]);
      mul(res, inv[sz[el]]);
    }
    mul(res, fact[n - 1]);
    printf("%d\\n", res);
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
vector<int> cntd;
vector<int> sz;
vector<int> cntu;
vector<int> par;
vector<vector<int>> adj;
vector<int> fact;
vector<int> inv;
vector<int> ans;
const int MOD = 1e9 + 7;

void mul(int &a, int b) { a = 1ll * a * b % MOD; }

void add(int &a, int b) {
  a += b;
  if (a >= MOD)
    a -= MOD;
}

ll pow1(ll n, ll p) {
  ll r = 1, c = n;
  while (p) {
    if (p & 1)
      (r *= c) %= MOD;
    p /= 2;
    (c *= c) %= MOD;
  }
  return r;
}

void dfs_d(int u, int p = -1) {
  sz[u] = 1;
  cntd[u] = 1;
  par[u] = p;
  for (auto el : adj[u]) {
    if (el == p)
      continue;
    dfs_d(el, u);
    sz[u] += sz[el];
    mul(cntd[u], cntd[el]);
    mul(cntd[u], inv[sz[el]]);
  }
  mul(cntd[u], fact[sz[u] - 1]);
}
int n;

void dfs_u(int u, int p = -1) {
  if (p == -1)
    cntu[u] = 1;
  else {
    cntu[u] = 1;
    int szp = n - sz[u];
    int szpp = szp - 1;
    for (auto el : adj[p]) {
      if (el == u || el == par[p])
        continue;
      mul(cntu[u], cntd[el]);
      mul(cntu[u], inv[sz[el]]);
      szpp -= sz[el];
    }
    mul(cntu[u], cntu[p]);
    mul(cntu[u], inv[szpp]);
    mul(cntu[u], fact[szp - 1]);
  }
  for (auto el : adj[u]) {
    if (el == p)
      continue;
    dfs_u(el, u);
  }
}

int main() {
  //  freopen("in.txt", "r", stdin);
  //  freopen("out.txt", "w", stdout);
  cin >> n;
  cntd.resize(n);
  sz.resize(n);
  cntu.resize(n);
  par.resize(n);
  adj.resize(n);
  fact.resize(n + 1);
  inv.resize(n + 1);
  ans.resize(n);
  fact[0] = inv[0] = 1;
  for (int i = 1; i <= n; i++) {
    fact[i] = 1ll * fact[i - 1] * i % MOD;
    inv[i] = pow1(fact[i], MOD - 2);
  }
  //  assert(1ll * fact[3] * inv[3] % MOD == 1);

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    scanf("%d %d", &u, &v);
    u--, v--;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  dfs_d(0);
  dfs_u(0);
  //  for(int i = 0;i < n;i++)
  //    cout << cntu[i] << endl;
  for (int i = 0; i < n; i++) {
    int szp = n - sz[i];
    int res = cntu[i];
    mul(res, inv[szp]);
    for (auto el : adj[i]) {
      if (el == par[i])
        continue;
      mul(res, cntd[el]);
      mul(res, inv[sz[el]]);
    }
    mul(res, fact[n - 1]);
    printf("%d\\n", res);
  }
  return 0;
}