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

#pragma GCC target("avx")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")

#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
const ll MOD = (ll)(1e9 + 7);
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (int)(n); (i)++)

template <typename T> T power(T x, long long n, const T &m) {
  if (n == 0)
    return 1;
  if (n == 1)
    return x;
  T tmp = power(x, n / 2, m);
  if (n % 2 == 0)
    return tmp * tmp % m;
  else
    return tmp * tmp % m * x % m;
}

template <typename T> T inverse(T x, T m) { return power(x, m - 2, m); }

int N;
int A[200000], B[200000];
vector<int> e[200000];

int fact[200001], inve[200001];

int sz[200000];
int par[200000];
vector<int> chi[200000];
int dfs(int n, int p) {
  par[n] = p;
  sz[n] = 1;
  rep(i, e[n].size()) {
    if (e[n][i] == p)
      continue;
    chi[n].pb(e[n][i]);
    sz[n] += dfs(e[n][i], n);
  }
  return sz[n];
}

ll dp[200000] = {};

ll calc(int n);

map<int, ll> memo[200000];
ll solve(int n, int b) {
  if (n == -1)
    return 1;
  if (b == par[n])
    return calc(n);
  else {
    if (memo[n].count(b) == 1)
      return memo[n][b];
    ll ret = calc(n);
    ret %= MOD;
    ret *= inve[sz[n] - 1];
    ret %= MOD;
    ret *= fact[N - sz[b] - 1];
    ret %= MOD;
    ret *= fact[sz[b]];
    ret %= MOD;
    ret *= inve[N - sz[n]];
    ret %= MOD;
    ret *= inverse(solve(b, n), MOD);
    ret %= MOD;
    ret *= solve(par[n], n);
    ret %= MOD;
    return memo[n][b] = ret;
  }
}

ll calc(int n) { // 親除く
  if (dp[n] != 0)
    return dp[n];
  ll ret = fact[sz[n] - 1];
  rep(i, chi[n].size()) {
    ret *= inve[sz[chi[n][i]]];
    ret %= MOD;
    ret *= calc(chi[n][i]);
    ret %= MOD;
  }
  return dp[n] = ret;
}

signed main() {
  scanf("%d", &N);
  rep(i, N - 1) {
    scanf("%d %d", A + i, B + i);
    A[i]--;
    B[i]--;
    e[A[i]].pb(B[i]);
    e[B[i]].pb(A[i]);
  }
  fact[0] = 1;
  inve[0] = 1;
  rep(i, N + 1) {
    fact[i + 1] = (fact[i] * (ll)(i + 1)) % MOD;
    inve[i + 1] = inverse((ll)fact[i + 1], MOD);
  }
  dfs(0, -1);
  rep(i, N) {
    ll ans = fact[N - 1];
    rep(j, chi[i].size()) {
      ans *= inve[sz[chi[i][j]]];
      ans %= MOD;
      ans *= calc(chi[i][j]);
      ans %= MOD;
    }
    if (i != 0) {
      ans *= inve[N - sz[i]];
      ans %= MOD;
      ans *= solve(par[i], i);
      ans %= MOD;
    }
    printf("%lld\\n", ans);
  }
}


Error: TLE

Fixed submission:

#pragma GCC target("avx")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")

#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
const ll MOD = (ll)(1e9 + 7);
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (int)(n); (i)++)

template <typename T> T power(T x, long long n, const T &m) {
  if (n == 0)
    return 1;
  if (n == 1)
    return x;
  T tmp = power(x, n / 2, m);
  if (n % 2 == 0)
    return tmp * tmp % m;
  else
    return tmp * tmp % m * x % m;
}

template <typename T> T inverse(T x, T m) { return power(x, m - 2, m); }

int N;
int A[200000], B[200000];
vector<int> e[200000];

int fact[200001], inve[200001];

int sz[200000];
int par[200000];
vector<int> chi[200000];
int dfs(int n, int p) {
  par[n] = p;
  sz[n] = 1;
  rep(i, e[n].size()) {
    if (e[n][i] == p)
      continue;
    chi[n].pb(e[n][i]);
    sz[n] += dfs(e[n][i], n);
  }
  return sz[n];
}

ll dp[200000] = {};

ll calc(int n);

ll solve(int n, int b) {
  if (n == -1)
    return 1;
  if (dp[n] == 0)
    calc(n);
  if (b == par[n])
    return dp[n];
  else {
    ll ret = dp[n];
    ret %= MOD;
    ret *= inve[sz[n] - 1];
    ret %= MOD;
    ret *= fact[N - sz[b] - 1];
    ret %= MOD;
    ret *= fact[sz[b]];
    ret %= MOD;
    ret *= inve[N - sz[n]];
    ret %= MOD;
    ret *= inverse(solve(b, n), MOD);
    ret %= MOD;
    ret *= solve(par[n], n);
    ret %= MOD;
    return ret;
  }
}

ll calc(int n) { // 親除く
  if (dp[n] != 0)
    return dp[n];
  ll ret = fact[sz[n] - 1];
  rep(i, chi[n].size()) {
    ret *= inve[sz[chi[n][i]]];
    ret %= MOD;
    ret *= calc(chi[n][i]);
    ret %= MOD;
  }
  return dp[n] = ret;
}

signed main() {
  scanf("%d", &N);
  rep(i, N - 1) {
    scanf("%d %d", A + i, B + i);
    A[i]--;
    B[i]--;
    e[A[i]].pb(B[i]);
    e[B[i]].pb(A[i]);
  }
  fact[0] = 1;
  inve[0] = 1;
  rep(i, N + 1) {
    fact[i + 1] = (fact[i] * (ll)(i + 1)) % MOD;
    inve[i + 1] = inverse((ll)fact[i + 1], MOD);
  }
  dfs(0, -1);
  rep(i, N) {
    ll ans = fact[N - 1];
    rep(j, chi[i].size()) {
      ans *= inve[sz[chi[i][j]]];
      ans %= MOD;
      ans *= calc(chi[i][j]);
      ans %= MOD;
    }
    if (i != 0) {
      ans *= inve[N - sz[i]];
      ans %= MOD;
      ans *= solve(par[i], i);
      ans %= MOD;
    }
    printf("%lld\\n", ans);
  }
}