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

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

const int N = 2e5 + 10, mod = 1e9 + 7;

int n, dp[N], sz[N], res[N], fact[N], invfact[N];
vector<int> adj[N];
vector<pair<int, int>> pref[N], suf[N];

int fp(int a, int b) {
  int ret = 1;
  while (b) {
    if (b & 1)
      ret = 1ll * ret * a % mod;
    b >>= 1;
    a = 1ll * a * a % mod;
  }
  return ret;
}

int com(int m, int r) {
  return 1ll * fact[m] * invfact[r] % mod * invfact[m - r] % mod;
}

int merge(int w1, int w2, int s1, int s2) {
  // cout << s1 + s2 << ' ' << s1 << ' ' << com(s1 + s2, s1) << '\\n';
  return 1ll * w1 * w2 % mod * com(s1 + s2, s1) % mod;
}

void dfs(int u, int pa) {
  dp[u] = 1;
  sz[u] = 0;
  for (int v : adj[u]) {
    if (v == pa)
      continue;
    dfs(v, u);
    dp[u] = 1ll * dp[u] * dp[v] % mod * com(sz[u] + sz[v], sz[u]) % mod;
    pref[u].push_back({dp[u], sz[u] + sz[v]});
    sz[u] += sz[v];
  }
  pair<int, int> t = {1, 0};
  int j = -1;
  for (int i = adj[u].size() - 1; i >= 0; i--) {
    int v = adj[u][i];
    if (v == pa) {
      j = i;
      continue;
    }
    t.first = merge(t.first, dp[v], t.second, sz[v]);
    t.second += sz[v];
    suf[u].push_back(t);
  }
  reverse(suf[u].begin(), suf[u].end());
  if (j > -1)
    adj[u].erase(adj[u].begin() + j);
  sz[u]++;
}

void dfs1(int u, int pa, int up = 1, int upsz = 0) {
  // cout << u << ' ' << up << ' ' << upsz << ' ' << dp[u] << '\\n';
  res[u] = merge(dp[u], up, sz[u] - 1, upsz);
  // if(u == 2) cout << adj[u].size() << '\\n';
  for (int i = 0; i < adj[u].size(); i++) {
    int v = adj[u][i];
    if (v == pa)
      continue;
    int Up = up, Upsz = upsz;
    if (i)
      Up = merge(Up, pref[u][i - 1].first, Upsz, pref[u][i - 1].second),
      Upsz += pref[u][i - 1].second;
    if (i < adj[u].size() - 1)
      Up = merge(Up, suf[u][i + 1].first, Upsz, suf[u][i + 1].second),
      Upsz += suf[u][i + 1].second;
    dfs1(v, u, Up, Upsz + 1);
  }
}

int main() {
  // ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  fact[0] = invfact[0] = 1;
  for (int i = 1; i < N; i++) {
    fact[i] = 1ll * fact[i - 1] * i % mod;
    invfact[i] = fp(fact[i], mod - 2);
  }
  int tc = 1;
  // scanf("%d", &tc);
  while (tc--) {
    scanf("%d", &n);
    for (int i = 0; i < n - 1; i++) {
      int u, v;
      scanf("%d%d", &u, &v);
      adj[u].push_back(v);
      adj[v].push_back(u);
    }

    dfs(1, 1);
    dfs1(1, 1);

    for (int i = 1; i <= n; i++)
      printf("%d\\n", res[i]);
  }
  return 0;
}

Error: TLE

Fixed submission:

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

const int N = 2e5 + 10, mod = 1e9 + 7;

map<int, int> C[N];
map<int, bool> vis[N];

int n, dp[N], sz[N], res[N];
vector<int> adj[N];
vector<pair<int, int>> pref[N], suf[N];

int com(int m, int r) {
  if (m < r || m < 0 || r < 0)
    return 0;
  if (m == r || r == 0)
    return 1;
  if (vis[m][r])
    return C[m][r];
  vis[m][r] = 1;
  return C[m][r] = (1ll * com(m - 1, r) + com(m - 1, r - 1)) % mod;
}

int merge(int w1, int w2, int s1, int s2) {
  // cout << s1 + s2 << ' ' << s1 << ' ' << com(s1 + s2, s1) << '\\n';
  return 1ll * w1 * w2 % mod * com(s1 + s2, s1) % mod;
}

void dfs(int u, int pa) {
  dp[u] = 1;
  sz[u] = 0;
  for (int v : adj[u]) {
    if (v == pa)
      continue;
    dfs(v, u);
    dp[u] = 1ll * dp[u] * dp[v] % mod * com(sz[u] + sz[v], sz[u]) % mod;
    pref[u].push_back({dp[u], sz[u] + sz[v]});
    sz[u] += sz[v];
  }
  pair<int, int> t = {1, 0};
  int j = -1;
  for (int i = adj[u].size() - 1; i >= 0; i--) {
    int v = adj[u][i];
    if (v == pa) {
      j = i;
      continue;
    }
    t.first = merge(t.first, dp[v], t.second, sz[v]);
    t.second += sz[v];
    suf[u].push_back(t);
  }
  reverse(suf[u].begin(), suf[u].end());
  if (j > -1)
    adj[u].erase(adj[u].begin() + j);
  sz[u]++;
}

void dfs1(int u, int pa, int up = 1, int upsz = 0) {
  // cout << u << ' ' << up << ' ' << upsz << ' ' << dp[u] << '\\n';
  res[u] = merge(dp[u], up, sz[u] - 1, upsz);
  // if(u == 2) cout << adj[u].size() << '\\n';
  for (int i = 0; i < adj[u].size(); i++) {
    int v = adj[u][i];
    if (v == pa)
      continue;
    int Up = up, Upsz = upsz;
    if (i)
      Up = merge(Up, pref[u][i - 1].first, Upsz, pref[u][i - 1].second),
      Upsz += pref[u][i - 1].second;
    if (i < adj[u].size() - 1)
      Up = merge(Up, suf[u][i + 1].first, Upsz, suf[u][i + 1].second),
      Upsz += suf[u][i + 1].second;
    dfs1(v, u, Up, Upsz + 1);
  }
}

int main() {
  // ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  int tc = 1;
  // scanf("%d", &tc);
  while (tc--) {
    scanf("%d", &n);
    for (int i = 0; i < n - 1; i++) {
      int u, v;
      scanf("%d%d", &u, &v);
      adj[u].push_back(v);
      adj[v].push_back(u);
    }

    dfs(1, 1);
    dfs1(1, 1);

    for (int i = 1; i <= n; i++)
      printf("%d\\n", res[i]);
  }
  return 0;
}