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

#include <algorithm>
#include <iostream>

#include <array>
#include <bitset>
#include <complex>
#include <cstring>
#include <deque>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <valarray>
#include <vector>

#include <cassert>
#include <chrono>
#include <cmath>
#include <functional>
#include <iomanip>
#include <numeric>
#include <random>

#include <memory>

using namespace std;
#define int long long

typedef long long ll;
typedef unsigned long long ull;
// typedef unsigned __int128 HASH;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ull, ull> pullull;
typedef pair<ll, int> plli;
typedef pair<double, int> pdi;
typedef pair<long double, int> pdbi;
typedef pair<int, pii> pipii;
typedef pair<ll, pll> plpll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<pii> vpii;
typedef vector<vector<int>> mat;

#define rep(i, n) for (int i = 0; i < (n); i++)
#define rep2(i, a, b) for (int i = (a); i < (b); i++)
#define rrep(i, n) for (int i = (n); i > 0; i--)
#define rrep2(i, a, b) for (int i = (a); i > b; i--)
#define fi first
#define se second
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()

template <class T>
using MINPQ = std::priority_queue<T, vector<T>, greater<T>>; // MINPQ<pii> pq;
template <class T> using MAXPQ = std::priority_queue<T>;

const ll hmod1 = 999999937;
const ll hmod2 = 1000000000 + 9;
const int INF = 1 << 30;
const ll INFLL = 1LL << 62;
const double EPS = 1e-12;
const ll mod = 1000000000 + 7;
// const ll mod = 998244353;
const int dx4[4] = {1, 0, -1, 0};
const int dy4[4] = {0, 1, 0, -1};
const int dx8[8] = {1, 1, 1, 0, 0, -1, -1, -1};
const int dy8[8] = {0, 1, -1, 1, -1, 0, 1, -1};
const long double pi = 3.141592653589793;

#define addm(X, Y) (X) = ((X) + ((Y) % mod) + mod) % mod
#define inside(y, x, h, w)                                                     \\
  (0 <= (y) && (y) < (h) && 0 <= (x) && (x) < (w)) ? true : false
#define stand(B, i)                                                            \\
  (((B >> i) & 1) == 1) ? true : false // Bのi番目のbitが1ならtrue, 0ならfalse

// debug
#define DEBUG

#define DUMPOUT cerr

#ifdef DEBUG
#define dump(...)                                                              \\
  DUMPOUT << #__VA_ARGS__ << " :[" << __FUNCTION__ << ":" << __LINE__ << "]"   \\
          << endl;                                                             \\
  DUMPOUT << "    ";                                                           \\
  dump_func(__VA_ARGS__)
#else
#define dump(...)
#endif

void dump_func() { DUMPOUT << endl; };

template <class Head, class... Tail>
void dump_func(Head &&head, Tail &&...tail) {
  DUMPOUT << head;
  if (sizeof...(Tail) == 0)
    DUMPOUT << " ";
  else
    DUMPOUT << ", ";
  dump_func(std::move(tail)...);
}

// ostream
template <typename T> ostream &operator<<(ostream &os, vector<T> &vec) {
  os << "[";
  for (int i = 0; i < vec.size(); i++)
    os << vec[i] << (i + 1 == vec.size() ? "" : ", ");
  os << "]";
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, pair<T, U> &pair_var) {
  os << "(" << pair_var.first << ", " << pair_var.second << ")";
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, map<T, U> &map_var) {
  os << "[";
  for (auto itr = map_var.begin(); itr != map_var.end(); itr++) {
    os << "(" << itr->first << ", " << itr->second << ")";
    itr++;
    if (itr != map_var.end())
      os << ", ";
    itr--;
  }
  os << "]";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, set<T> &set_var) {
  os << "[";
  for (auto itr = set_var.begin(); itr != set_var.end(); itr++) {
    os << *itr;
    ++itr;
    if (itr != set_var.end())
      os << ", ";
    itr--;
  }
  os << "]";
  return os;
}

int n;
vector<vector<int>> g;
int cnt[200000 + 5];
vector<vector<int>> sub(200000 + 5);

const int COMMAX = 2000000 + 5; // nCkのnの上限
vector<long long> fact(COMMAX), finv(COMMAX), inv(COMMAX);

long long powmod(long long n, long long k) {
  if (k < 0)
    return 0;
  if (k == 0)
    return 1;
  else if (k % 2 == 0) {
    long long tmp = powmod(n, k / 2);
    return tmp * tmp % mod;
  } else
    return powmod(n, k - 1) * n % mod;
}

void init_com() {
  fact[0] = fact[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < COMMAX; i++) {
    fact[i] = fact[i - 1] * i % mod;
    inv[i] = mod - inv[mod % i] * (mod / i) % mod;
    finv[i] = finv[i - 1] * inv[i] % mod;
  }
}

long long inverse(long long n) { return powmod(n, mod - 2); }

long long nCk(int n, int k) {
  if (n < k || n < 0 || k < 0)
    return 0;
  return fact[n] * (finv[k] * finv[n - k] % mod) % mod;
}

int calc_sub(int now, int par) {
  int m = 0;
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int x = calc_sub(nx, now);
    sub[now].push_back(x);
    m += x;
  }
  sub[now].push_back(1);
  return m + 1;
}

int dfs(int now, int par) {
  int tmp = 1;
  vector<int> sub_sz;
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int x = dfs(nx, now);
    tmp = (tmp * x) % mod;
  }
  int m = 0;
  rep(i, sub[now].size()) m += sub[now][i];
  m--;
  rep(i, sub[now].size() - 1LL) {
    // cerr << tmp << " " << nCk(m, sub[now][i]) << " " << m << " " <<
    // sub[now][i] << endl;
    tmp = (tmp * nCk(m, sub[now][i])) % mod;
    m -= sub[now][i];
  }
  cnt[now] = tmp;
  // cerr << now << " " << cnt[now] << " " << tmp << " " << sub[now] <<endl;
  return cnt[now];
}

int ans[200000 + 5];

void dfs1(int now, int par, int d) {
  ans[now] = cnt[now];
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int nx_sub = 0;
    rep(i, sub[nx].size()) nx_sub += sub[nx][i];
    sub[nx].push_back(d - nx_sub);
    int sum = nx_sub + d - nx_sub;
    sum--;
    int tmp = 1;
    bool self = true;
    for (auto x : sub[nx]) {
      if (x == 1 && self) {
        self = false;
        continue;
      }
      tmp = tmp * nCk(sum, x) % mod;
      sum -= x;
    }
    int pre = cnt[now];
    cnt[now] = cnt[now] * inverse(cnt[nx]) % mod;
    cnt[now] = cnt[now] * inverse(nCk(d - 1, nx_sub)) % mod;
    for (auto nx1 : g[nx]) {
      tmp = tmp * cnt[nx1] % mod;
    }
    int pre1 = cnt[nx];
    cnt[nx] = tmp;
    dfs1(nx, now, nx_sub + d - nx_sub);
    cnt[now] = pre;
    cnt[nx] = pre1;
    sub[nx].pop_back();
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> n;
  init_com();
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  calc_sub(0, -1);
  dfs(0, -1);
  /*
  rep(i, n) {
    cerr << i << " " << sub[i] << endl;
    cerr << i << " " << cnt[i] << endl;
  }
  */
  int d = 0;
  rep(i, sub[0].size()) d += sub[0][i];
  dfs1(0, -1, d);
  rep(i, n) { cout << ans[i] << endl; }
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>

#include <array>
#include <bitset>
#include <complex>
#include <cstring>
#include <deque>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <valarray>
#include <vector>

#include <cassert>
#include <chrono>
#include <cmath>
#include <functional>
#include <iomanip>
#include <numeric>
#include <random>

#include <memory>

using namespace std;
#define int long long

typedef long long ll;
typedef unsigned long long ull;
// typedef unsigned __int128 HASH;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ull, ull> pullull;
typedef pair<ll, int> plli;
typedef pair<double, int> pdi;
typedef pair<long double, int> pdbi;
typedef pair<int, pii> pipii;
typedef pair<ll, pll> plpll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<pii> vpii;
typedef vector<vector<int>> mat;

#define rep(i, n) for (int i = 0; i < (n); i++)
#define rep2(i, a, b) for (int i = (a); i < (b); i++)
#define rrep(i, n) for (int i = (n); i > 0; i--)
#define rrep2(i, a, b) for (int i = (a); i > b; i--)
#define fi first
#define se second
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()

template <class T>
using MINPQ = std::priority_queue<T, vector<T>, greater<T>>; // MINPQ<pii> pq;
template <class T> using MAXPQ = std::priority_queue<T>;

const ll hmod1 = 999999937;
const ll hmod2 = 1000000000 + 9;
const int INF = 1 << 30;
const ll INFLL = 1LL << 62;
const double EPS = 1e-12;
const ll mod = 1000000000 + 7;
// const ll mod = 998244353;
const int dx4[4] = {1, 0, -1, 0};
const int dy4[4] = {0, 1, 0, -1};
const int dx8[8] = {1, 1, 1, 0, 0, -1, -1, -1};
const int dy8[8] = {0, 1, -1, 1, -1, 0, 1, -1};
const long double pi = 3.141592653589793;

#define addm(X, Y) (X) = ((X) + ((Y) % mod) + mod) % mod
#define inside(y, x, h, w)                                                     \\
  (0 <= (y) && (y) < (h) && 0 <= (x) && (x) < (w)) ? true : false
#define stand(B, i)                                                            \\
  (((B >> i) & 1) == 1) ? true : false // Bのi番目のbitが1ならtrue, 0ならfalse

// debug
#define DEBUG

#define DUMPOUT cerr

#ifdef DEBUG
#define dump(...)                                                              \\
  DUMPOUT << #__VA_ARGS__ << " :[" << __FUNCTION__ << ":" << __LINE__ << "]"   \\
          << endl;                                                             \\
  DUMPOUT << "    ";                                                           \\
  dump_func(__VA_ARGS__)
#else
#define dump(...)
#endif

void dump_func() { DUMPOUT << endl; };

template <class Head, class... Tail>
void dump_func(Head &&head, Tail &&...tail) {
  DUMPOUT << head;
  if (sizeof...(Tail) == 0)
    DUMPOUT << " ";
  else
    DUMPOUT << ", ";
  dump_func(std::move(tail)...);
}

// ostream
template <typename T> ostream &operator<<(ostream &os, vector<T> &vec) {
  os << "[";
  for (int i = 0; i < vec.size(); i++)
    os << vec[i] << (i + 1 == vec.size() ? "" : ", ");
  os << "]";
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, pair<T, U> &pair_var) {
  os << "(" << pair_var.first << ", " << pair_var.second << ")";
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, map<T, U> &map_var) {
  os << "[";
  for (auto itr = map_var.begin(); itr != map_var.end(); itr++) {
    os << "(" << itr->first << ", " << itr->second << ")";
    itr++;
    if (itr != map_var.end())
      os << ", ";
    itr--;
  }
  os << "]";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, set<T> &set_var) {
  os << "[";
  for (auto itr = set_var.begin(); itr != set_var.end(); itr++) {
    os << *itr;
    ++itr;
    if (itr != set_var.end())
      os << ", ";
    itr--;
  }
  os << "]";
  return os;
}

int n;
vector<vector<int>> g;
int cnt[200000 + 5];
vector<vector<int>> sub(200000 + 5);

const int COMMAX = 2000000 + 5; // nCkのnの上限
vector<long long> fact(COMMAX), finv(COMMAX), inv(COMMAX);

long long powmod(long long n, long long k) {
  if (k < 0)
    return 0;
  if (k == 0)
    return 1;
  else if (k % 2 == 0) {
    long long tmp = powmod(n, k / 2);
    return tmp * tmp % mod;
  } else
    return powmod(n, k - 1) * n % mod;
}

void init_com() {
  fact[0] = fact[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < COMMAX; i++) {
    fact[i] = fact[i - 1] * i % mod;
    inv[i] = mod - inv[mod % i] * (mod / i) % mod;
    finv[i] = finv[i - 1] * inv[i] % mod;
  }
}

long long nCk(int n, int k) {
  if (n < k || n < 0 || k < 0)
    return 0;
  return fact[n] * (finv[k] * finv[n - k] % mod) % mod;
}

int calc_sub(int now, int par) {
  int m = 0;
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int x = calc_sub(nx, now);
    sub[now].push_back(x);
    m += x;
  }
  return m + 1;
}

int dfs(int now, int par) {
  int tmp = 1;
  vector<int> sub_sz;
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int x = dfs(nx, now);
    tmp = (tmp * x) % mod;
  }
  int m = 0;
  rep(i, sub[now].size()) m += sub[now][i];
  rep(i, sub[now].size()) {
    // cerr << tmp << " " << nCk(m, sub[now][i]) << " " << m << " " <<
    // sub[now][i] << endl;
    tmp = (tmp * nCk(m, sub[now][i])) % mod;
    m -= sub[now][i];
  }
  cnt[now] = tmp;
  // cerr << now << " " << cnt[now] << " " << tmp << " " << sub[now] <<endl;
  return cnt[now];
}

int ans[200000 + 5];

void dfs1(int now, int par) {
  ans[now] = cnt[now];
  for (auto nx : g[now]) {
    if (nx == par)
      continue;
    int tmp1 = cnt[now];
    int tmp2 = cnt[nx];
    int m = 0;
    // sub[now], sub[nx] update
    rep(i, sub[nx].size()) { m += sub[nx][i]; }
    m++;
    int tmp_sub1 = m;
    sub[now].erase(find(sub[now].begin(), sub[now].end(), m));
    m = 0;
    rep(i, sub[now].size()) m += sub[now][i];
    sub[nx].push_back(m + 1);
    // int tmp_sub2 = m + 1;
    // cnt[now] update
    // m = 0;
    // rep(i, sub[now].size()) m += sub[now][i];
    int tmp = 1;
    rep(i, sub[now].size()) {
      tmp = tmp * nCk(m, sub[now][i]) % mod;
      m -= sub[now][i];
    }
    for (auto nx1 : g[now]) {
      if (nx1 == nx)
        continue;
      tmp = (tmp * cnt[nx1]) % mod;
    }
    cnt[now] = tmp;
    // cnt[nx] update
    m = 0;
    rep(i, sub[nx].size()) m += sub[nx][i];
    tmp = 1;
    rep(i, sub[nx].size()) {
      tmp = tmp * nCk(m, sub[nx][i]) % mod;
      m -= sub[nx][i];
    }
    for (auto nx2 : g[nx]) {
      tmp = (tmp * cnt[nx2]) % mod;
    }
    cnt[nx] = tmp;
    dfs1(nx, now);
    cnt[now] = tmp1;
    cnt[nx] = tmp2;
    sub[now].push_back(tmp_sub1);
    // sub[nx].erase(find(sub[nx].begin(), sub[nx].end(), tmp_sub2));
    sub[nx].pop_back();
    // sub[nx] = tmp_sub2;
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> n;
  init_com();
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  calc_sub(0, -1);
  dfs(0, -1);
  /*
  rep(i, n) {
    cerr << i << " " << sub[i] << endl;
    cerr << i << " " << cnt[i] << endl;
  }
  */
  dfs1(0, -1);
  rep(i, n) { cout << ans[i] << endl; }
}