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

#include <algorithm>
#include <iostream>
#include <queue>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

const ll MOD = 1000000007;

ll modpow(ll x, ll n) {
  ll res = 1;
  ll mul = x;
  while (n) {
    if (n % 2)
      res = res * mul % MOD;
    mul = mul * mul % MOD;
    n /= 2;
  }
  return res;
}

struct Combination {
  int maxn;
  vector<ll> fmemo;
  vector<ll> imemo;
  vector<ll> ifmemo;

  Combination(int maxn) : maxn(maxn) {
    fmemo.resize(maxn + 1);
    imemo.resize(maxn + 1);
    ifmemo.resize(maxn + 1);
    for (ll i = 1; i <= maxn; i++) {
      imemo[i] = modpow(i, MOD - 2);
    }
  }

  ll fact(ll n) {
    if (n == 0)
      return 1;
    if (fmemo[n])
      return fmemo[n];
    return fmemo[n] = n * fact(n - 1) % MOD;
  }

  ll ifact(ll n) {
    if (n == 0)
      return 1;
    if (ifmemo[n])
      return ifmemo[n];
    return ifmemo[n] = imemo[n] * ifact(n - 1) % MOD;
  }

  ll com(ll n, ll k) {
    if (k > n)
      return 0;
    if (k == n || k == 0)
      return 1;
    return (fact(n) * ifact(k) % MOD) * ifact(n - k) % MOD;
  }
};

struct Graph {
  int n;
  vector<vector<int>> g;

  Graph(int n) : n(n) { g.resize(n); }

  void init(int n_) {
    n = n_;
    g.resize(n_);
  }

  void add_edge(int from, int to) { g[from].push_back(to); }
};

struct Tree // create tree(directed) from graph(undirected)
{
  int n;
  int root;
  vector<vector<int>> t;
  vector<int> par;
  vector<int> dpt;

  void init(Graph &g, int root_) {
    n = g.n;
    root = root_;
    t.resize(n);
    par.resize(n);
    dpt.resize(n);
    fill(dpt.begin(), dpt.end(), -1);
    queue<int> que;
    par[root] = -1;
    dpt[root] = 0;
    que.push(root);
    while (que.size()) {
      int pa = que.front();
      que.pop();
      for (int ch : g.g[pa]) {
        if (dpt[ch] == -1) {
          t[pa].push_back(ch);
          par[ch] = pa;
          dpt[ch] = dpt[pa] + 1;
          que.push(ch);
        }
      }
    }
  }

  Tree() {}

  Tree(Graph &g, int root_) { init(g, root_); }
};

P p[200005];
ll dp[200005];
ll c[200005];
ll ans[200005];

int main() {
  int n;
  cin >> n;
  Combination com(n);
  Graph g(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g.add_edge(a, b);
    g.add_edge(b, a);
  }
  Tree t(g, 0);
  for (int i = 0; i < n; i++)
    p[i] = P(t.dpt[i], i);
  sort(p, p + n, greater<P>());
  fill(dp, dp + n, 1);
  for (int i = 0; i < n; i++) {
    int u = p[i].second;
    for (int v : t.t[u]) {
      dp[u] = dp[u] * dp[v] % MOD;
      dp[u] = dp[u] * com.ifact(c[v] + 1) % MOD;
      c[u] += c[v] + 1;
    }
    dp[u] = dp[u] * com.fact(c[u]) % MOD;
  }
  fill(ans, ans + n, 1);
  for (int i = 0; i < n; i++)
    ans[i] = dp[i];
  sort(p, p + n);
  for (int i = 1; i < n; i++) {
    int u = p[i].second;
    int w = t.par[u];
    ans[u] = ans[u] * ans[w] % MOD * modpow(dp[u], MOD - 2) % MOD * (c[u] + 1) %
             MOD * modpow(n - c[u] - 1, MOD - 2) % MOD;
  }
  for (int i = 0; i < n; i++)
    cout << ans[i] << endl;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <queue>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

const ll MOD = 1000000007;

ll modpow(ll x, ll n) {
  ll res = 1;
  ll mul = x;
  while (n) {
    if (n % 2)
      res = res * mul % MOD;
    mul = mul * mul % MOD;
    n /= 2;
  }
  return res;
}

struct Combination {
  int maxn;
  vector<ll> fmemo;
  vector<ll> imemo;
  vector<ll> ifmemo;

  Combination(int maxn) : maxn(maxn) {
    fmemo.resize(maxn + 1);
    imemo.resize(maxn + 1);
    ifmemo.resize(maxn + 1);
    for (ll i = 1; i <= maxn; i++) {
      imemo[i] = modpow(i, MOD - 2);
    }
  }

  ll fact(ll n) {
    if (n == 0)
      return 1;
    if (fmemo[n])
      return fmemo[n];
    return fmemo[n] = n * fact(n - 1) % MOD;
  }

  ll ifact(ll n) {
    if (n == 0)
      return 1;
    if (ifmemo[n])
      return ifmemo[n];
    return ifmemo[n] = imemo[n] * ifact(n - 1) % MOD;
  }

  ll com(ll n, ll k) {
    if (k > n)
      return 0;
    if (k == n || k == 0)
      return 1;
    return (fact(n) * ifact(k) % MOD) * ifact(n - k) % MOD;
  }
};

struct Graph {
  int n;
  vector<vector<int>> g;

  Graph(int n) : n(n) { g.resize(n); }

  void init(int n_) {
    n = n_;
    g.resize(n_);
  }

  void add_edge(int from, int to) { g[from].push_back(to); }
};

struct Tree // create tree(directed) from graph(undirected)
{
  int n;
  int root;
  vector<vector<int>> t;
  vector<int> par;
  vector<int> dpt;

  void init(Graph &g, int root_) {
    n = g.n;
    root = root_;
    t.resize(n);
    par.resize(n);
    dpt.resize(n);
    fill(dpt.begin(), dpt.end(), -1);
    queue<int> que;
    par[root] = -1;
    dpt[root] = 0;
    que.push(root);
    while (que.size()) {
      int pa = que.front();
      que.pop();
      for (int ch : g.g[pa]) {
        if (dpt[ch] == -1) {
          t[pa].push_back(ch);
          par[ch] = pa;
          dpt[ch] = dpt[pa] + 1;
          que.push(ch);
        }
      }
    }
  }

  Tree() {}

  Tree(Graph &g, int root_) { init(g, root_); }
};

P p[200005];
ll dp[200005];
ll c[200005];
ll pa[200005];
ll pc[200005];
ll ans[200005];

int main() {
  int n;
  cin >> n;
  Combination com(n);
  Graph g(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g.add_edge(a, b);
    g.add_edge(b, a);
  }
  Tree t(g, 0);
  for (int i = 0; i < n; i++)
    p[i] = P(t.dpt[i], i);
  sort(p, p + n, greater<P>());
  fill(dp, dp + n, 1);
  for (int i = 0; i < n; i++) {
    int u = p[i].second;
    for (int v : t.t[u]) {
      dp[u] = dp[u] * dp[v] % MOD;
      dp[u] = dp[u] * com.ifact(c[v] + 1) % MOD;
      c[u] += c[v] + 1;
    }
    dp[u] = dp[u] * com.fact(c[u]) % MOD;
  }
  fill(pa, pa + n, 1);
  fill(ans, ans + n, 1);
  ans[0] = dp[0];
  sort(p, p + n);
  for (int i = 1; i < n; i++) {
    int u = p[i].second;
    int w = t.par[u];
    pa[u] = pa[u] * pa[w] % MOD;
    pa[u] = pa[u] * com.ifact(n - (c[w] + 1)) % MOD;
    for (int v : t.t[w]) {
      if (v == u)
        continue;
      pa[u] = pa[u] * dp[v] % MOD;
      pa[u] = pa[u] * com.ifact(c[v] + 1) % MOD;
    }
    pa[u] = pa[u] * com.fact(n - (c[u] + 1) - 1) % MOD;
    ans[u] = ans[u] * pa[u] % MOD;
    ans[u] = ans[u] * com.ifact(n - (c[u] + 1)) % MOD;
    for (int v : t.t[u]) {
      ans[u] = ans[u] * dp[v] % MOD;
      ans[u] = ans[u] * com.ifact(c[v] + 1) % MOD;
    }
    ans[u] = ans[u] * com.fact(n - 1) % MOD;
  }
  for (int i = 0; i < n; i++)
    cout << ans[i] << endl;
}