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

// clang-format off
#include <bits/stdc++.h>

#define mp make_pair
#define fst first
#define snd second
#define forn(i,n) for (int i = 0; i < int(n); i++)
#define forn1(i,n) for (int i = 1; i <= int(n); i++)
#define popcnt __builtin_popcount
#define ffs __builtin_ffs
#define ctz __builtin_ctz
#define clz __builtin_clz
#define all(a) (a).begin(), (a).end()
 
using namespace std;

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
using pii = pair<int,int>;
using pli = pair<ll,int>;
using pil = pair<int,ll>;
using pll = pair<ll,ll>;
template <typename T> using vec = vector<T>;
using vi = vec<int>;
using vl = vec<ll>;
template <typename T> using que = queue<T>;
template <typename T> using deq = deque<T>;

template <typename T> T id(T b) {return b;};
template <typename T> void chmax(T &x, T y) {if (x < y) x = y;}
template <typename T> void chmin(T &x, T y) {if (x > y) x = y;}
template <typename S, typename K> bool contains(S &s, K k) { return s.find(k) != s.end(); }
void fastio() { ios_base::sync_with_stdio(false); cin.tie(nullptr); }
constexpr ll TEN(int n) { if (n == 0) return 1LL; else return 10LL*TEN(n-1); }
// clang-format on

const ll MOD = TEN(9) + 7;
const int MAX_N = 200000;
int n;
vi g[MAX_N];
ll cnt[MAX_N];
ll dp[MAX_N];
ll fact[MAX_N + 1];
ll fact_inv[MAX_N + 1];
vl ans;

template <long long mod> long long mod_pow(long long a, long long b) {
  if (b == 0)
    return 1;
  long long ret = mod_pow<mod>(a, b / 2);
  ret = ret * ret % mod;
  if (b % 2 == 1)
    ret = a * ret % mod;
  return ret;
}

void dfs(int i, int p) {
  ll ans = 1;
  ll c = 0;
  for (int j : g[i]) {
    if (j == p)
      continue;
    dfs(j, i);
    ans = ans * dp[j] % MOD * fact_inv[cnt[j]] % MOD;
    c = (c + cnt[j]) % MOD;
  }

  ans = ans * fact[c] % MOD;
  cnt[i] = c + 1;
  dp[i] = ans;
}

void pre_calc() {
  fact[0] = 1;
  for (ll i = 1; i <= MAX_N; i++) {
    fact[i] = i * fact[i - 1] % MOD;
  }

  for (ll i = 0; i <= MAX_N; i++) {
    fact_inv[i] = mod_pow<MOD>(fact[i], MOD - 2) % MOD;
  }
}

ll inv(ll x) { return mod_pow<MOD>(x, MOD - 2); }

void go(int i, int p) {
  ans[i] = dp[i];

  for (int j : g[i]) {
    if (j == p)
      continue;
    ll dpi = dp[i];
    ll ci = cnt[i];
    ll dpj = dp[j];
    ll cj = cnt[j];

    dp[i] = dp[i] * fact_inv[n - 1] % MOD * inv(dp[j]) % MOD * fact[cj] % MOD *
            fact[n - 1 - cj] % MOD;
    cnt[i] = n - cj;
    // printf("dp[%d]=%lld cnt[%d]=%lld\\n", i, dp[i], i, cnt[i]);
    dp[j] = dp[j] * fact_inv[cj - 1] % MOD * dp[i] % MOD * fact_inv[cnt[i]] %
            MOD * fact[n - 1] % MOD;
    cnt[j] = n;

    go(j, i);
    dp[i] = dpi;
    dp[j] = dpj;
    cnt[i] = ci;
    cnt[j] = cj;
  }
}

int main() {
  fastio();

  pre_calc();

  cin >> n;
  forn(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }

  ans = vl(n);
  dfs(0, -1);
  go(0, -1);

  for (ll x : ans) {
    cout << x << "\\n";
  }

  return 0;
}

Error: TLE

Fixed submission:

// clang-format off
#include <bits/stdc++.h>

#define mp make_pair
#define fst first
#define snd second
#define forn(i,n) for (int i = 0; i < int(n); i++)
#define forn1(i,n) for (int i = 1; i <= int(n); i++)
#define popcnt __builtin_popcount
#define ffs __builtin_ffs
#define ctz __builtin_ctz
#define clz __builtin_clz
#define all(a) (a).begin(), (a).end()
 
using namespace std;

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
using pii = pair<int,int>;
using pli = pair<ll,int>;
using pil = pair<int,ll>;
using pll = pair<ll,ll>;
template <typename T> using vec = vector<T>;
using vi = vec<int>;
using vl = vec<ll>;
template <typename T> using que = queue<T>;
template <typename T> using deq = deque<T>;

template <typename T> T id(T b) {return b;};
template <typename T> void chmax(T &x, T y) {if (x < y) x = y;}
template <typename T> void chmin(T &x, T y) {if (x > y) x = y;}
template <typename S, typename K> bool contains(S &s, K k) { return s.find(k) != s.end(); }
void fastio() { ios_base::sync_with_stdio(false); cin.tie(nullptr); }
constexpr ll TEN(int n) { if (n == 0) return 1LL; else return 10LL*TEN(n-1); }
// clang-format on

const ll MOD = TEN(9) + 7;
const int MAX_N = 200000;
int n;
vi g[MAX_N];
ll cnt[MAX_N];
ll dp[MAX_N];
ll fact[MAX_N + 1];
ll inv[MAX_N + 1];
vl ans;

template <long long mod> long long mod_pow(long long a, long long b) {
  if (b == 0)
    return 1;
  long long ret = mod_pow<mod>(a, b / 2);
  ret = ret * ret % mod;
  if (b % 2 == 1)
    ret = a * ret % mod;
  return ret;
}

void dfs(int i, int p) {
  ll ans = 1;
  ll c = 0;
  for (int j : g[i]) {
    if (j == p)
      continue;
    dfs(j, i);
    ans = ans * dp[j] % MOD * inv[cnt[j]] % MOD;
    c = (c + cnt[j]) % MOD;
  }

  ans = ans * fact[c] % MOD;
  cnt[i] = c + 1;
  dp[i] = ans;
}

void pre_calc() {
  fact[0] = 1;
  for (ll i = 1; i <= MAX_N; i++) {
    fact[i] = i * fact[i - 1] % MOD;
  }

  for (ll i = 0; i <= MAX_N; i++) {
    inv[i] = mod_pow<MOD>(fact[i], MOD - 2) % MOD;
  }
}

inline void calc(int i, int j = -1) {
  ll c = 0;
  ll ans = 1;
  for (int k : g[i]) {
    if (k == j)
      continue;
    ans = ans * dp[k] % MOD * inv[cnt[k]] % MOD;
    c = (c + cnt[k]) % MOD;
  }

  ans = ans * fact[c] % MOD;
  cnt[i] = c + 1;
  dp[i] = ans;
}

inline void go(int i, int p) {
  ans[i] = dp[i];

  for (int j : g[i]) {
    if (j == p)
      continue;
    ll dpi = dp[i];
    ll ci = cnt[i];
    ll dpj = dp[j];
    ll cj = cnt[j];
    calc(i, j);
    calc(j);
    go(j, i);
    dp[i] = dpi;
    dp[j] = dpj;
    cnt[i] = ci;
    cnt[j] = cj;
  }
}

int main() {
  fastio();

  pre_calc();

  cin >> n;
  forn(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }

  ans = vl(n);
  dfs(0, -1);
  go(0, -1);

  for (ll x : ans) {
    cout << x << "\\n";
  }

  return 0;
}