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

#include <bits/stdc++.h>
using namespace std;

const int64_t MOD = 1e9 + 7;
class ModInt {
public:
  int64_t x;
  ModInt(int64_t x = 0) : x((x % MOD + MOD) % MOD) {}
  ModInt operator-() const { return ModInt(-x); }
  ModInt &operator+=(const ModInt a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  ModInt operator+(const ModInt a) const { return ModInt(*this) += a; }
  ModInt operator-(const ModInt a) const { return ModInt(*this) -= a; }
  ModInt operator*(const ModInt a) const { return ModInt(*this) *= a; }
  ModInt pow(int64_t n) const {
    if (n == 0)
      return 1;
    ModInt t = pow(n >> 1);
    t *= t;
    if (n & 1)
      t *= *this;
    return t;
  }
  ModInt inv() const { return pow(MOD - 2); }
  ModInt &operator/=(const ModInt a) { return *this *= a.inv(); }
  ModInt operator/(const ModInt a) const { return ModInt(*this) /= a; }
};
istream &operator>>(istream &is, const ModInt &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const ModInt &a) { return os << a.x; }

vector<ModInt> fac;
vector<ModInt> fac_inv;
vector<ModInt> inv;
void init_comb(int mx) {
  fac = vector<ModInt>(mx + 1), fac_inv = vector<ModInt>(mx + 1),
  inv = vector<ModInt>(mx + 1);
  fac.at(0) = fac.at(1) = 1;
  fac_inv.at(0) = fac_inv.at(1) = 1;
  inv.at(1) = 1;
  for (int i = 2; i <= mx; i++) {
    fac.at(i) = fac.at(i - 1) * i;
    inv.at(i) = -inv.at(MOD % i) * (MOD / i);
    fac_inv.at(i) = fac_inv.at(i - 1) * inv.at(i);
  }
}
ModInt comb(int n, int k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac.at(n) * fac_inv.at(k) * fac_inv.at(n - k);
}

class Data {
public:
  ModInt c;
  int s;
  Data(ModInt c = 1, int s = 0) : c(c), s(s) {}
  Data &operator+=(const Data a) {
    s += a.s;
    c *= comb(s, a.s) * a.c;
    return *this;
  }
  Data operator-(const Data a) const {
    return Data(c / (comb(s, a.s) * a.c), s - a.s);
  }
  Data addRoot() { return Data(c, s + 1); }
};

vector<vector<int>> tree;
vector<Data> dp;

void dfs(int v, int p = -1) {
  Data data = Data();
  for (auto u : tree.at(v)) {
    if (u == p)
      continue;
    dfs(u, v);
    data += dp.at(u).addRoot();
  }
  dp.at(v) = data;
}

int main() {
  int N;
  cin >> N;
  vector<int> a(N - 1), b(N - 1);
  for (int i = 0; i < N - 1; i++) {
    cin >> a.at(i) >> b.at(i);
  }

  tree = vector<vector<int>>(N);
  for (int i = 0; i < N - 1; i++) {
    int u = a.at(i) - 1, v = b.at(i) - 1;
    tree.at(u).push_back(v);
    tree.at(v).push_back(u);
  }

  init_comb(N);

  dp = vector<Data>(N);
  dfs(0);

  queue<pair<int, int>> q;
  for (auto v : tree.at(0)) {
    q.push(make_pair(v, 0));
  }
  while (!q.empty()) {
    int v = q.front().first, p = q.front().second;
    q.pop();
    dp.at(v) += (dp.at(p) - dp.at(v).addRoot()).addRoot();

    for (auto u : tree.at(v)) {
      if (u == p)
        continue;
      q.push(make_pair(u, v));
    }
  }

  for (int i = 0; i < N; i++) {
    cout << dp.at(i).c << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

const int64_t MOD = 1e9 + 7;
class ModInt {
public:
  int64_t x;
  ModInt(int64_t x = 0) : x((x % MOD + MOD) % MOD) {}
  ModInt operator-() const { return ModInt(-x); }
  ModInt &operator+=(const ModInt a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator-=(const ModInt a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  ModInt &operator*=(const ModInt a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  ModInt operator+(const ModInt a) const { return ModInt(*this) += a; }
  ModInt operator-(const ModInt a) const { return ModInt(*this) -= a; }
  ModInt operator*(const ModInt a) const { return ModInt(*this) *= a; }
  ModInt pow(int64_t n) const {
    if (n == 0)
      return 1;
    ModInt t = pow(n >> 1);
    t *= t;
    if (n & 1)
      t *= *this;
    return t;
  }
  ModInt inv() const { return pow(MOD - 2); }
  ModInt &operator/=(const ModInt a) { return *this *= a.inv(); }
  ModInt operator/(const ModInt a) const { return ModInt(*this) /= a; }
};
istream &operator>>(istream &is, const ModInt &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const ModInt &a) { return os << a.x; }

vector<ModInt> fac;
vector<ModInt> fac_inv;
vector<ModInt> inv;
void init_comb(int mx) {
  fac = vector<ModInt>(mx + 1), fac_inv = vector<ModInt>(mx + 1),
  inv = vector<ModInt>(mx + 1);
  fac.at(0) = fac.at(1) = 1;
  fac_inv.at(0) = fac_inv.at(1) = 1;
  inv.at(1) = 1;
  for (int i = 2; i <= mx; i++) {
    fac.at(i) = fac.at(i - 1) * i;
    inv.at(i) = -inv.at(MOD % i) * (MOD / i);
    fac_inv.at(i) = fac_inv.at(i - 1) * inv.at(i);
  }
}
ModInt comb(int n, int k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac.at(n) * fac_inv.at(k) * fac_inv.at(n - k);
}

vector<vector<int>> tree;
vector<vector<ModInt>> c;
vector<vector<int>> s;

pair<ModInt, int> dfs(int v, int p = -1) {
  ModInt res_v_c = 1;
  int res_v_s = 0;
  for (int i = 0; i < (int)tree.at(v).size(); i++) {
    int u = tree.at(v).at(i);
    if (u == p)
      continue;
    if (s.at(v).at(i) < 0) {
      pair<ModInt, int> res_u = dfs(u, v);
      c.at(v).at(i) = res_u.first;
      s.at(v).at(i) = res_u.second;
    }
    res_v_s += s.at(v).at(i);
    res_v_c = comb(res_v_s, s.at(v).at(i)) * res_v_c * c.at(v).at(i);
  }
  res_v_s += 1;
  return make_pair(res_v_c, res_v_s);
}

int main() {
  int N;
  cin >> N;
  vector<int> a(N - 1), b(N - 1);
  for (int i = 0; i < N - 1; i++) {
    cin >> a.at(i) >> b.at(i);
  }

  tree = vector<vector<int>>(N);
  for (int i = 0; i < N - 1; i++) {
    int u = a.at(i) - 1, v = b.at(i) - 1;
    tree.at(u).push_back(v);
    tree.at(v).push_back(u);
  }

  init_comb(N);

  c = vector<vector<ModInt>>(N);
  s = vector<vector<int>>(N);
  for (int i = 0; i < N; i++) {
    c.at(i) = vector<ModInt>(tree.at(i).size());
    s.at(i) = vector<int>(tree.at(i).size(), -1);
  }

  vector<ModInt> res(N);
  vector<bool> visited(N, false);
  queue<int> q;
  q.push(0);
  while (!q.empty()) {
    int v = q.front();
    q.pop();
    res.at(v) = dfs(v).first;
    visited.at(v) = true;

    for (auto u : tree.at(v)) {
      if (visited.at(u))
        continue;
      q.push(u);
    }
  }

  for (int i = 0; i < N; i++) {
    cout << res.at(i) << endl;
  }

  return 0;
}