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

#include <bits/stdc++.h>
#define int long long
#define pii pair<int, int>

using namespace std;
const int inf = 1e9 + 7;
const int max_n = 2e5 + 7;
vector<int> gr[max_n];
map<int, int> dp[max_n], sz[max_n];
int fact[max_n];

inline int mul(int a, int b) { return a * b % inf; }

inline int power(int a, int b) {
  int res = 1;
  while (b) {
    if (b & 1)
      res = mul(res, a);

    a = mul(a, a);

    b >>= 1;
  }

  return res;
}

inline int inv(int a) { return power(a, inf - 2); }

inline int C(int n, int k) {
  return mul(fact[n], inv(mul(fact[k], fact[n - k])));
}

int iter = 0;

set<pii> temp;

void fuck(int v) {
  int c_sz = 0;
  int res = 1;
  for (int j : gr[v]) {
    c_sz += sz[j][v];
    res = mul(res, mul(dp[j][v], C(c_sz, sz[j][v])));
  }

  // cout << v << ":" << '\\n';

  // cout << c_sz << " " << res << '\\n';

  for (int j : gr[v]) {
    dp[v][j] = mul(res, inv(mul(dp[j][v], C(c_sz, sz[j][v]))));
    sz[v][j] = c_sz - sz[j][v] + 1;

    // cout << j << " " << dp[v][j] << " " << sz[v][j] << " " << dp[j][v] << " "
    // << sz[j][v] << '\\n';
  }

  // cout << '\\n' << '\\n';
}

void dfs(int v, int pr) {
  if (dp[v].size() >= 2) {
    fuck(v);
    return;
  }

  iter++;
  if (iter > 2 * max_n) {
    int aa = 0;
    fact[0] /= aa;
  }
  // temp.emplace(v, pr);

  sz[v][pr] = 0;
  dp[v][pr] = 1;

  for (int i : gr[v]) {
    if (i == pr)
      continue;
    if (dp[i].find(v) == dp[i].end())
      dfs(i, v);
    sz[v][pr] += sz[i][v];
    dp[v][pr] = mul(dp[v][pr], mul(dp[i][v], C(sz[v][pr], sz[i][v])));
  }
  sz[v][pr]++;
}

inline void solve() {
  fact[0] = 1;
  for (int i = 1; i < max_n; i++)
    fact[i] = mul(fact[i - 1], i);

  int n;
  cin >> n;
  for (int i = 0; i < n - 1; i++) {
    int v1, v2;
    cin >> v1 >> v2;
    v1--;
    v2--;
    gr[v1].emplace_back(v2);
    gr[v2].emplace_back(v1);
  }

  // cout << "scan" << endl;

  for (int i = 0; i < n; i++) {
    for (int pr : gr[i]) {
      if (dp[i].find(pr) == dp[i].end())
        dfs(i, pr);
    }
  }

  // cout << "dfs" << endl;

  for (int i = 0; i < n; i++) {
    int c_sz = 0;
    int res = 1;
    for (int j : gr[i]) {
      c_sz += sz[j][i];
      res = mul(res, mul(dp[j][i], C(c_sz, sz[j][i])));
    }

    cout << res << '\\n';
  }
}

signed main() {
  /*freopen("input.txt", "r", stdin);
freopen("output.txt", "w", stdout);*/

  cin.tie(0);
  cout.tie(0);
  ios_base::sync_with_stdio(0);

  solve();
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define int long long
#define pii pair<int, int>

using namespace std;
const int inf = 1e9 + 7;
const int max_n = 2e5 + 7;
vector<int> gr[max_n];
map<int, int> dp[max_n], sz[max_n];
int fact[max_n];

inline int mul(int a, int b) { return a * b % inf; }

inline int power(int a, int b) {
  int res = 1;
  while (b) {
    if (b & 1)
      res = mul(res, a);

    a = mul(a, a);

    b >>= 1;
  }

  return res;
}

inline int inv(int a) { return power(a, inf - 2); }

inline int C(int n, int k) {
  return mul(fact[n], inv(mul(fact[k], fact[n - k])));
}

int iter = 0;

set<pii> temp;

void dfs(int v, int pr) {
  iter++;
  if (iter > 2 * max_n) {
    int aa = 0;
    fact[0] /= aa;
  }
  // temp.emplace(v, pr);

  sz[v][pr] = 0;
  dp[v][pr] = 1;

  for (int i : gr[v]) {
    if (i == pr)
      continue;
    if (dp[i].find(v) == dp[i].end())
      dfs(i, v);
    sz[v][pr] += sz[i][v];
    dp[v][pr] = mul(dp[v][pr], mul(dp[i][v], C(sz[v][pr], sz[i][v])));
  }
  sz[v][pr]++;
}

inline void solve() {
  fact[0] = 1;
  for (int i = 1; i < max_n; i++)
    fact[i] = mul(fact[i - 1], i);

  int n;
  cin >> n;
  for (int i = 0; i < n - 1; i++) {
    int v1, v2;
    cin >> v1 >> v2;
    v1--;
    v2--;
    gr[v1].emplace_back(v2);
    gr[v2].emplace_back(v1);
  }

  // cout << "scan" << endl;

  for (int i = 0; i < n; i++) {
    for (int pr : gr[i]) {
      if (dp[i].find(pr) == dp[i].end())
        dfs(i, pr);
    }
  }

  // cout << "dfs" << endl;

  for (int i = 0; i < n; i++) {
    int c_sz = 0;
    int res = 1;
    for (int j : gr[i]) {
      c_sz += sz[j][i];
      res = mul(res, mul(dp[j][i], C(c_sz, sz[j][i])));
    }

    cout << res << '\\n';
  }
}

signed main() {
  /*freopen("input.txt", "r", stdin);
freopen("output.txt", "w", stdout);
*/
  cin.tie(0);
  cout.tie(0);
  ios_base::sync_with_stdio(0);

  solve();
}