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

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \\
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

static const long long mod = 1000000007;
struct mint {
  long long x;
  mint(ll x = 0) : x(x % mod) {}
  mint &operator+=(const mint a) {
    (x += a.x) %= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    (x += mod - a.x) %= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint ret(*this);
    return ret += a;
  }
  mint operator-(const mint a) const {
    mint ret(*this);
    return ret -= a;
  }
  mint operator*(const mint a) const {
    mint ret(*this);
    return ret *= a;
  }
  mint pow(ll t) const {
    if (t == 0)
      return mint(1);
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return (*this) *= a.inv(); }
  mint operator/(const mint a) const {
    mint ret(*this);
    return ret /= a;
  }
};
ostream &operator<<(ostream &os, const mint &x) {
  os << x.x;
  return os;
}
class modutils {
  vector<mint> fact, invfact;

public:
  modutils(int n = 200005) : fact(n + 1), invfact(n + 1) {
    fact[0] = 1;
    for (int i = 1; i <= n; i++)
      fact[i] = fact[i - 1] * i;
    invfact[n] = fact[n].inv();
    for (int i = n; i >= 1; i--)
      invfact[i - 1] = invfact[i] * i;
  }
  mint pow(mint x, ll n) { return x.pow(n); }
  mint comb(ll n, ll k) {
    if (n < 0 || k < 0 || n < k)
      return 0;
    return fact[n] * invfact[k] * invfact[n - k];
  }
  mint perm(ll n, ll k) {
    if (n < 0 || k < 0 || n < k)
      return 0;
    return fact[n] * invfact[n - k];
  }
  mint fac(ll n) { return fact[n]; }
};
modutils mutil;

struct ST {
  int src, dst;
  int cn;
  mint cnt;
  int rev;
};

int N;
vector<ST> G[200005];

void solve(int v, int p, ST &e) {
  vector<int> cns;
  mint ret(1);

  rep(i, G[v].size()) {
    if (G[v][i].dst == p)
      continue;

    if (G[v][i].cn == -1) {
      solve(G[v][i].dst, v, G[v][i]);
    }

    cns.push_back(G[v][i].cn);
    ret *= G[v][i].cnt;
  }

  if (cns.size() == 0) {
    e.cn = 1;
    e.cnt = ret;
  } else {
    int sum = 0;
    rep(i, cns.size()) { sum += cns[i]; }
    mint tmp = mutil.fac(sum);
    rep(i, cns.size()) { tmp /= mutil.fac(cns[i]); }
    e.cn = sum + 1;
    e.cnt = ret * tmp;
  }
}

void solve2() {
  queue<pair<int, int>> que;
  que.push(make_pair(0, -1));

  while (!que.empty()) {
    pair<int, int> pr = que.front();
    que.pop();
    int v = pr.first;
    int p = pr.second;

    int sum = 0;
    mint ret(1);
    mint pp(1);
    rep(i, G[v].size()) {
      sum += G[v][i].cn;
      ret *= G[v][i].cnt;
      pp *= mutil.fac(G[v][i].cn);
    }

    rep(i, G[v].size()) {
      int tsum = sum - G[v][i].cn;
      mint tret = ret / G[v][i].cnt;
      mint ttmp = mutil.fac(tsum);
      ttmp /= pp / mutil.fac(G[v][i].cn);

      G[G[v][i].dst][G[v][i].rev].cn = tsum + 1;
      G[G[v][i].dst][G[v][i].rev].cnt = tret * ttmp;
    }

    rep(i, G[v].size()) {
      if (G[v][i].dst == p)
        continue;
      que.push(make_pair(G[v][i].dst, v));
    }
  }
}

int main() {
  cin >> N;
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;

    G[a].push_back((ST){a, b, -1, 0, (int)(G[b].size())});
    G[b].push_back((ST){b, a, -1, 0, (int)(G[a].size() - 1)});
  }

  {
    ST st = (ST){-1, 0, -1, 0, -1};
    solve(0, -1, st);
    solve2();
  }

  rep(i, N) {
    ST st = (ST){-1, i, -1, 0, -1};
    solve(i, -1, st);
    cout << st.cnt << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \\
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

static const long long mod = 1000000007;
struct mint {
  long long x;
  mint(ll x = 0) : x(x % mod) {}
  mint &operator+=(const mint a) {
    (x += a.x) %= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    (x += mod - a.x) %= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint ret(*this);
    return ret += a;
  }
  mint operator-(const mint a) const {
    mint ret(*this);
    return ret -= a;
  }
  mint operator*(const mint a) const {
    mint ret(*this);
    return ret *= a;
  }
  mint pow(ll t) const {
    if (t == 0)
      return mint(1);
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return (*this) *= a.inv(); }
  mint operator/(const mint a) const {
    mint ret(*this);
    return ret /= a;
  }
};
ostream &operator<<(ostream &os, const mint &x) {
  os << x.x;
  return os;
}
class modutils {
  vector<mint> fact, invfact;

public:
  modutils(int n = 200005) : fact(n + 1), invfact(n + 1) {
    fact[0] = 1;
    for (int i = 1; i <= n; i++)
      fact[i] = fact[i - 1] * i;
    invfact[n] = fact[n].inv();
    for (int i = n; i >= 1; i--)
      invfact[i - 1] = invfact[i] * i;
  }
  mint pow(mint x, ll n) { return x.pow(n); }
  mint comb(ll n, ll k) {
    if (n < 0 || k < 0 || n < k)
      return 0;
    return fact[n] * invfact[k] * invfact[n - k];
  }
  mint perm(ll n, ll k) {
    if (n < 0 || k < 0 || n < k)
      return 0;
    return fact[n] * invfact[n - k];
  }
  mint fac(ll n) { return fact[n]; }
};
modutils mutil;

struct ST {
  int src, dst;
  int cn;
  mint cnt;
  int rev;
};

int N;
vector<ST> G[200005];

void solve(int v, int p, ST &e) {
  vector<int> cns;
  mint ret(1);
  rep(i, G[v].size()) {
    if (G[v][i].dst == p)
      continue;

    if (G[v][i].cn == -1) {
      solve(G[v][i].dst, v, G[v][i]);
    }

    cns.push_back(G[v][i].cn);
    ret *= G[v][i].cnt;
  }

  if (cns.size() == 0) {
    e.cn = 1;
    e.cnt = ret;
  } else {
    int sum = 0;
    rep(i, cns.size()) { sum += cns[i]; }
    mint tmp = mutil.fac(sum);
    rep(i, cns.size()) { tmp /= mutil.fac(cns[i]); }
    e.cn = sum + 1;
    e.cnt = ret * tmp;
  }
}

int main() {
  cin >> N;
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;

    G[a].push_back((ST){a, b, -1, 0, (int)(G[b].size())});
    G[b].push_back((ST){b, a, -1, 0, (int)(G[a].size() - 1)});
  }

  rep(i, N) {
    ST st = (ST){-1, i, -1, 0, -1};
    solve(i, -1, st);
    cout << st.cnt << endl;
  }

  return 0;
}