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

#include <bits/stdc++.h>

using namespace std;

using pii = pair<int, int>;
using ll = long long;
#define rep(i, j) for (int i = 0; i < (int)(j); i++)
#define repeat(i, j, k) for (int i = (j); i < (int)(k); i++)
#define all(v) v.begin(), v.end()
#define debug(x) cerr << #x << " : " << x << endl

template <class T> bool set_min(T &a, const T &b) {
  return a > b ? a = b, true : false;
}
template <class T> bool set_max(T &a, const T &b) {
  return a < b ? a = b, true : false;
}
// vector
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &a : v)
    is >> a;
  return is;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (const T &t : v)
    os << "\\t" << t;
  return os << endl;
}
// pair
template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &v) {
  return os << "<" << v.first << ", " << v.second << ">";
}

const int INF = 1 << 30;
const ll INFL = 1LL << 60;
const ll MOD = 1000000007LL;

class ModInt {
public:
  static const ll MOD = 1000000007LL;
  ll x;
  ModInt() : x(0){};
  ModInt(ll x) {
    while (x < 0)
      x += MOD;
    x %= MOD;
    this->x = x;
  }
  ModInt &operator+=(const ModInt &l) {
    x += l.x;
    if (x >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt &l) {
    x -= l.x;
    if (x < 0)
      x += MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt &l) {
    x = (x * l.x) % MOD;
    return *this;
  }
  ModInt operator+(const ModInt &l) const { return ModInt(x + l.x); }
  ModInt operator-(const ModInt &l) const { return ModInt(x - l.x); }
  ModInt operator*(const ModInt &l) const { return ModInt(x * l.x); }
};

// @warning rの部分はmodとってはいけない
ModInt pow(const ModInt &n, ll r) {
  if (r == 0)
    return ModInt(1);
  ModInt ret = pow(n, r / 2);
  ret *= ret;
  if (r % 2 != 0)
    ret = ret * n;
  return ret;
}

// @waring nはMODと互いに素
ModInt inverse(const ModInt &n) { return pow(n, ModInt::MOD - 2); }

ModInt factorial(const ModInt &n) {
  ll ret = 1;
  for (ll i = n.x; i > 1; i--) {
    ret *= i;
  }
  return ModInt(ret);
}

ostream &operator<<(std::ostream &os, const ModInt &a) {
  os << a.x;
  return os;
}
istream &operator>>(std::istream &is, ModInt &a) {
  is >> a.x;
  return is;
}

vector<ll> memo_size;
vector<ll> memo_dp1;
vector<ll> memo_dp2;
vector<ll> factorials;

int calc_size(int now, int prev, vector<vector<int>> &E) {
  if (memo_size[now] >= 0)
    return memo_size[now];
  int res = 1;
  for (int c : E[now])
    if (c != prev)
      res += calc_size(c, now, E);
  memo_size[now] = res;
  return res;
}

ll calc_routes(int now, int prev, vector<vector<int>> &E) {
  if (memo_dp1[now] >= 0)
    return memo_dp1[now];
  vector<ll> sizes;
  vector<ll> cnts;
  for (int c : E[now]) {
    if (c != prev) {
      sizes.push_back(calc_size(c, now, E));
      cnts.push_back(calc_routes(c, now, E));
    }
  }
  if (sizes.size() == 0) {
    memo_dp1[now] = 1;
    return 1;
  } else if (sizes.size() == 1) {
    memo_dp1[now] = cnts[0];
    return cnts[0];
  }
  ModInt a = factorials[accumulate(all(sizes), 0ll)];
  for (ll c : cnts)
    a *= c;
  ModInt b = 1;
  for (ll s : sizes)
    b *= factorials[s];
  memo_dp1[now] = (a * inverse(b)).x;
  return memo_dp1[now];
}

void calc_routes2(int now, int prev, ll parent_size, ll parent_val,
                  vector<vector<int>> &E) {
  // cerr << "At " << now << " from " << prev << " | p size " << parent_size <<
  // " p val " << parent_val << endl;
  vector<ll> sizes;
  vector<ll> cnts;
  vector<ll> cs;
  for (int c : E[now]) {
    if (c != prev) {
      sizes.push_back(calc_size(c, now, E));
      cnts.push_back(calc_routes(c, now, E));
      cs.push_back(c);
    }
  }
  if (sizes.size() == 0) {
    memo_dp2[now] = parent_val;
    return;
  }
  ll total_size = accumulate(all(sizes), 0ll) + parent_size;
  ModInt a = 1;
  for (ll c : cnts)
    a *= c;
  ModInt b = 1;
  for (ll s : sizes)
    b *= factorials[s];
  // self:
  {
    auto aa = a * factorials[total_size] * parent_val;
    auto bb = b * factorials[parent_size];
    memo_dp2[now] = (aa * inverse(bb)).x;
  }
  rep(i, sizes.size()) {
    auto aa =
        a * inverse(cnts[i]) * factorials[total_size - sizes[i]] * parent_val;
    auto bb = b * inverse(factorials[sizes[i]]) * factorials[parent_size];
    calc_routes2(cs[i], now, total_size - sizes[i] + 1, (aa * inverse(bb)).x,
                 E);
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int N;
  cin >> N;
  vector<vector<int>> E(N);
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    E[a].push_back(b);
    E[b].push_back(a);
  }
  memo_size = vector<ll>(N, -1);
  memo_dp1 = vector<ll>(N, -1);
  memo_dp2 = vector<ll>(N, -1);
  factorials.resize(N + 10);
  factorials[0] = 1;
  repeat(i, 1, N + 10) factorials[i] = (factorials[i - 1] * i) % MOD;

  calc_routes(0, -1, E);
  calc_routes2(0, -1, 0, 1, E);
  rep(i, N) cout << memo_dp2[i] << endl;
  // debug(memo_dp1);
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

using pii = pair<int, int>;
using ll = long long;
#define rep(i, j) for (int i = 0; i < (int)(j); i++)
#define repeat(i, j, k) for (int i = (j); i < (int)(k); i++)
#define all(v) v.begin(), v.end()
#define debug(x) cerr << #x << " : " << x << endl

template <class T> bool set_min(T &a, const T &b) {
  return a > b ? a = b, true : false;
}
template <class T> bool set_max(T &a, const T &b) {
  return a < b ? a = b, true : false;
}
// vector
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &a : v)
    is >> a;
  return is;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (const T &t : v)
    os << "\\t" << t;
  return os << endl;
}
// pair
template <class T, class U>
ostream &operator<<(ostream &os, const pair<T, U> &v) {
  return os << "<" << v.first << ", " << v.second << ">";
}

const int INF = 1 << 30;
const ll INFL = 1LL << 60;
const ll MOD = 1000000007LL;

class ModInt {
public:
  static const ll MOD = 1000000007LL;
  ll x;
  ModInt() : x(0){};
  ModInt(ll x) {
    while (x < 0)
      x += MOD;
    x %= MOD;
    this->x = x;
  }
  ModInt &operator+=(const ModInt &l) {
    x += l.x;
    if (x >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt &l) {
    x -= l.x;
    if (x < 0)
      x += MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt &l) {
    x = (x * l.x) % MOD;
    return *this;
  }
  ModInt operator+(const ModInt &l) const { return ModInt(x + l.x); }
  ModInt operator-(const ModInt &l) const { return ModInt(x - l.x); }
  ModInt operator*(const ModInt &l) const { return ModInt(x * l.x); }
};

// @warning rの部分はmodとってはいけない
ModInt pow(const ModInt &n, ll r) {
  if (r == 0)
    return ModInt(1);
  ModInt ret = pow(n, r / 2);
  ret *= ret;
  if (r % 2 != 0)
    ret = ret * n;
  return ret;
}

// @waring nはMODと互いに素
ModInt inverse(const ModInt &n) { return pow(n, ModInt::MOD - 2); }

ModInt factorial(const ModInt &n) {
  ll ret = 1;
  for (ll i = n.x; i > 1; i--) {
    ret *= i;
  }
  return ModInt(ret);
}

ostream &operator<<(std::ostream &os, const ModInt &a) {
  os << a.x;
  return os;
}
istream &operator>>(std::istream &is, ModInt &a) {
  is >> a.x;
  return is;
}

vector<ll> memo_size;
vector<ll> memo_dp1;
vector<ll> memo_dp2;
vector<ll> factorials;

int calc_size(int now, int prev, vector<vector<int>> &E) {
  if (memo_size[now] >= 0)
    return memo_size[now];
  int res = 1;
  for (int c : E[now])
    if (c != prev)
      res += calc_size(c, now, E);
  memo_size[now] = res;
  return res;
}

ll calc_routes(int now, int prev, vector<vector<int>> &E) {
  if (memo_dp1[now] >= 0)
    return memo_dp1[now];
  vector<ll> sizes;
  vector<ll> cnts;
  for (int c : E[now]) {
    if (c != prev) {
      sizes.push_back(calc_size(c, now, E));
      cnts.push_back(calc_routes(c, now, E));
    }
  }
  if (sizes.size() == 0) {
    memo_dp1[now] = 1;
    return 1;
  } else if (sizes.size() == 1) {
    memo_dp1[now] = cnts[0];
    return cnts[0];
  }
  ll val = cnts[0];
  ll size = sizes[0];
  for (int i = 1; i < sizes.size(); i++) {
    ll nxt_size = size + sizes[i];
    val = (ModInt(factorials[nxt_size]) *
           inverse(ModInt(factorials[size] * factorials[sizes[i]])) * val *
           cnts[i])
              .x;
    size = nxt_size;
  }
  memo_dp1[now] = val;
  return val;
}

void calc_routes2(int now, int prev, ll parent_size, ll parent_val,
                  vector<vector<int>> &E) {
  // cerr << "At " << now << " from " << prev << " | p size " << parent_size <<
  // " p val " << parent_val << endl;
  vector<ll> sizes;
  vector<ll> cnts;
  vector<ll> cs;
  for (int c : E[now]) {
    if (c != prev) {
      sizes.push_back(calc_size(c, now, E));
      cnts.push_back(calc_routes(c, now, E));
      cs.push_back(c);
    }
  }
  if (sizes.size() == 0) {
    memo_dp2[now] = parent_val;
    return;
  }
  ll val = cnts[0];
  ll size = sizes[0];
  for (int i = 1; i < sizes.size(); i++) {
    ll nxt_size = size + sizes[i];
    val = (ModInt(factorials[nxt_size]) *
           inverse(ModInt(factorials[size] * factorials[sizes[i]])) * val *
           cnts[i])
              .x;
    size = nxt_size;
  }
  // self:
  {
    ll nxt_size = size + parent_size;
    val = (ModInt(factorials[nxt_size]) *
           inverse(ModInt(factorials[size] * factorials[parent_size])) * val *
           parent_val)
              .x;
    memo_dp2[now] = val;
    size = nxt_size;
  }
  for (int j = 0; j < sizes.size(); j++) {
    ll val = 1;
    ll size = 0;
    for (int i = 0; i < sizes.size(); i++) {
      if (i == j)
        continue;
      ll nxt_size = size + sizes[i];
      val = (ModInt(factorials[nxt_size]) *
             inverse(ModInt(factorials[size] * factorials[sizes[i]])) * val *
             cnts[i])
                .x;
      size = nxt_size;
    }
    {
      ll nxt_size = size + parent_size;
      val = (ModInt(factorials[nxt_size]) *
             inverse(ModInt(factorials[size] * factorials[parent_size])) * val *
             parent_val)
                .x;
      size = nxt_size;
    }
    calc_routes2(cs[j], now, size + 1, val, E);
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int N;
  cin >> N;
  vector<vector<int>> E(N);
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    E[a].push_back(b);
    E[b].push_back(a);
  }
  memo_size = vector<ll>(N, -1);
  memo_dp1 = vector<ll>(N, -1);
  memo_dp2 = vector<ll>(N, -1);
  factorials.resize(N + 10);
  factorials[0] = 1;
  repeat(i, 1, N + 10) factorials[i] = (factorials[i - 1] * i) % MOD;

  calc_routes(0, -1, E);
  calc_routes2(0, -1, 0, 1, E);
  rep(i, N) cout << memo_dp2[i] << endl;
  // debug(memo_dp1);
  return 0;
}