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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iterator>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define LINF2 1223300000000000000
#define INF 2140000000
#define INF2 1140000000
const long long MOD = 1000000007;
// const long long MOD = 998244353;

using namespace std;

const int max_comb = 210000;
vector<ll> fac(max_comb + 1);  // n! (mod M)
vector<ll> ifac(max_comb + 1); // k!^(-1) (mod M)

ll mpow(ll x, ll n) { // x^n(mod M)
  ll ans = 1;
  while (n != 0) {
    if (n & 1)
      ans = ans * x % MOD;
    x = x * x % MOD;
    n = n >> 1;
  }
  return ans;
}

ll minv(ll x) { return mpow(x, MOD - 2); }

ll comb(int a, int b) { // C(a,b) = a! * b!^(-1) * (a-b)^(-1)
  if (a == 0 && b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  ll tmp = ifac[a - b] * ifac[b] % MOD;
  return tmp * fac[a] % MOD;
}

ll perm(int a, int b) { // P(a,b) = a! * (a-b)!^(-1)
  if (b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  ll tmp = ifac[a - b] % MOD;
  return tmp * fac[a] % MOD;
}

void pre_comb() {
  fac[0] = 1;
  ifac[0] = 1;
  for (int i = 0; i < max_comb; i++) {
    fac[i + 1] = fac[i] * (i + 1) % MOD;       // n!(mod M)
    ifac[i + 1] = ifac[i] * minv(i + 1) % MOD; // k!^(-1) (mod M)
  }
  return;
}

vector<vector<int>> g;
vector<int> num0;
vector<ll> ans0;
vector<ll> ans00;

int dfs(int par, int curr, ll *pret) {
  int siz = (int)g[curr].size();
  vector<int> v0;
  int i;
  int sum = 0;
  ll ans = 1;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    ll ret = 0;
    int num = dfs(curr, next, &ret);
    sum += num;
    v0.push_back(num);
    ans = ans * ret % MOD;
  }

  int siz0 = (int)v0.size();
  int num1 = sum;
  for (i = 0; i < siz0; i++) {
    ans = ans * comb(num1, v0[i]) % MOD;
    num1 -= v0[i];
  }

  num0[curr] = sum + 1;
  ans0[curr] = ans;
  (*pret) = ans;
  return sum + 1;
}

void dfs2(int par, int curr, int num, ll ret0) {
  int siz = (int)g[curr].size();
  vector<int> v0;
  int i;
  int sum = num;
  v0.push_back(num);
  ll ans = ret0;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    v0.push_back(num0[next]);
    sum += num0[next];
    ans = ans * ans0[next] % MOD;
  }
  ll ansall = ans;
  {
    int siz0 = (int)v0.size();
    int num1 = sum;

    for (i = 0; i < siz0; i++) {
      ansall = ansall * comb(num1, v0[i]) % MOD;
      num1 -= v0[i];
    }
    ans00[curr] = ansall;
  }

  int cnt = 1;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    {
      int siz0 = (int)v0.size();
      int num1 = sum - num0[next];
      ll anstmp =
          ansall * minv(ans0[next]) % MOD * minv(comb(sum, num0[next])) % MOD;
      dfs2(curr, next, num1 + 1, anstmp);
    }
    cnt++;
  }

  return;
}

void solve() {
  int n;
  scanf("%d", &n);
  g.resize(n);
  ans0.resize(n);
  num0.resize(n);
  ans00.resize(n);
  int i;
  for (i = 0; i < n - 1; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  pre_comb();
  ll ret = 0;
  int num = dfs(-1, 0, &ret);
  assert(num == n);

  dfs2(-1, 0, 0, 1);

  for (i = 0; i < n; i++) {
    printf("%lld\\n", ans00[i]);
  }

  return;
}

int main(int argc, char *argv[]) {
#if 1
  solve();
#else
  int T;
  scanf("%d", &T);
  while (T--) {
    solve();
  }
#endif
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iterator>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define LINF2 1223300000000000000
#define INF 2140000000
#define INF2 1140000000
const long long MOD = 1000000007;
// const long long MOD = 998244353;

using namespace std;

const int max_comb = 210000;
vector<ll> fac(max_comb + 1);  // n! (mod M)
vector<ll> ifac(max_comb + 1); // k!^(-1) (mod M)

ll mpow(ll x, ll n) { // x^n(mod M)
  ll ans = 1;
  while (n != 0) {
    if (n & 1)
      ans = ans * x % MOD;
    x = x * x % MOD;
    n = n >> 1;
  }
  return ans;
}

ll minv(ll x) { return mpow(x, MOD - 2); }

ll comb(int a, int b) { // C(a,b) = a! * b!^(-1) * (a-b)^(-1)
  if (a == 0 && b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  ll tmp = ifac[a - b] * ifac[b] % MOD;
  return tmp * fac[a] % MOD;
}

ll perm(int a, int b) { // P(a,b) = a! * (a-b)!^(-1)
  if (b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  ll tmp = ifac[a - b] % MOD;
  return tmp * fac[a] % MOD;
}

void pre_comb() {
  fac[0] = 1;
  ifac[0] = 1;
  for (int i = 0; i < max_comb; i++) {
    fac[i + 1] = fac[i] * (i + 1) % MOD;       // n!(mod M)
    ifac[i + 1] = ifac[i] * minv(i + 1) % MOD; // k!^(-1) (mod M)
  }
  return;
}

vector<vector<int>> g;
vector<int> num0;
vector<ll> ans0;
vector<ll> ans00;

int dfs(int par, int curr, ll *pret) {
  int siz = (int)g[curr].size();
  vector<int> v0;
  int i;
  int sum = 0;
  ll ans = 1;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    ll ret = 0;
    int num = dfs(curr, next, &ret);
    sum += num;
    v0.push_back(num);
    ans = ans * ret % MOD;
  }

  int siz0 = (int)v0.size();
  int num1 = sum;
  for (i = 0; i < siz0; i++) {
    ans = ans * comb(num1, v0[i]) % MOD;
    num1 -= v0[i];
  }

  num0[curr] = sum + 1;
  ans0[curr] = ans;
  (*pret) = ans;
  return sum + 1;
}

void dfs2(int par, int curr, int num, ll ret0) {
  int siz = (int)g[curr].size();
  vector<int> v0;
  int i;
  int sum = num;
  v0.push_back(num);
  ll ans = ret0;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    v0.push_back(num0[next]);
    sum += num0[next];
    ans = ans * ans0[next] % MOD;
  }
  {
    int siz0 = (int)v0.size();
    int num1 = sum;
    ll ans5 = ans;
    for (i = 0; i < siz0; i++) {
      ans5 = ans5 * comb(num1, v0[i]) % MOD;
      num1 -= v0[i];
    }
    ans00[curr] = ans5;
  }

  int cnt = 1;
  for (i = 0; i < siz; i++) {
    int next = g[curr][i];
    if (par == next)
      continue;

    {
      int siz0 = (int)v0.size();
      int num1 = sum - num0[next];
      ll ans5 = ans * minv(ans0[next]) % MOD;
      int k;
      for (k = 0; k < siz0; k++) {
        if (k == cnt)
          continue;
        ans5 = ans5 * comb(num1, v0[k]) % MOD;
        num1 -= v0[k];
      }
      dfs2(curr, next, sum - num0[next] + 1, ans5);
    }
    cnt++;
  }

  return;
}

void solve() {
  int n;
  scanf("%d", &n);
  g.resize(n);
  ans0.resize(n);
  num0.resize(n);
  ans00.resize(n);
  int i;
  for (i = 0; i < n - 1; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  pre_comb();
  ll ret = 0;
  int num = dfs(-1, 0, &ret);
  assert(num == n);

  dfs2(-1, 0, 0, 1);

  for (i = 0; i < n; i++) {
    printf("%lld\\n", ans00[i]);
  }

  return;
}

int main(int argc, char *argv[]) {
#if 1
  solve();
#else
  int T;
  scanf("%d", &T);
  while (T--) {
    solve();
  }
#endif
  return 0;
}