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

#include <bits/stdc++.h>

using namespace std;
using ll = long long int;
constexpr ll M = 1e9 + 7;
inline ll sum_1(const ll n) { return ((n * (n + 1)) / 2) % M; }
inline ll sum_2(const ll n) { return ((n * (n + 1) * (2 * n + 1)) / 6) % M; }
inline ll sum_3(const ll n) {
  return (((n * (n + 1)) / 2) % M * ((n * (n + 1)) / 2) % M) % M;
}
inline ll sum_1(const ll i, const ll j) {
  return (sum_1(j) - sum_1(i - 1) + M) % M;
}
inline ll sum_2(const ll i, const ll j) {
  return (sum_2(j) - sum_2(i - 1) + M) % M;
}
inline ll sum_3(const ll i, const ll j) {
  return (sum_3(j) - sum_3(i - 1) + M) % M;
}

ll pow_f(ll b, ll e) {
  ll ans = 1;
  b = b % M;
  while (e) {
    if (e & 1) {
      ans = (ans * b) % M;
    }
    b = (b * b) % M;
    e >>= 1;
  }
  return ans;
}

ll nCr(ll n, ll r) {
  if (r > n || n < 0 || r < 0) {
    return 0;
  }
  if (n == r || r == 0) {
    return 1;
  }
  r = min(r, n - r);
  ll ans = 1;
  for (ll i = 1; i <= r; ++i) {
    ans = (ans * (n - i + 1)) / (i);
  }
  return ans;
}

ll nCr(tuple<ll, ll> t) {
  auto n = get<0>(t);
  auto k = get<1>(t);
  return nCr(n + k - 1, k);
}

unordered_map<ll, ll> fact;
ll factorial(ll n) {
  if (fact.find(n) != fact.end()) {
    return fact[n];
  }
  ll ans = 1;
  for (ll i = 1; i <= n; ++i) {
    ans = (ans * i) % M;
  }
  fact[n] = ans;
  return ans;
}

unordered_map<ll, ll> modinvmap;
inline ll modinv(ll n) {
  if (modinvmap.find(n) != modinvmap.end()) {
    return modinvmap[n];
  }
  auto ans = pow_f(n, M - 2);
  modinvmap[n] = ans;
  return ans;
}

class Tree {
  vector<vector<int>> adj;
  struct info {
    ll tp; // no. of toposorts
    ll sz; // size of subtree (including the root)
  };

  vector<bool> visited;
  vector<info> dp;
  vector<ll> ans;

public:
  Tree(int n) : visited(n + 1, false), adj(n + 1), dp(n + 1), ans(n + 1, -1) {
    // Precalc Factorial
    ll fact_t = 1;
    for (ll i = 1; i <= 2e5 + 5; ++i) {
      fact_t = (fact_t * i) % M;
      fact[i] = fact_t;
    }
  }

  void addEdge(int u, int v) {
    adj[u].push_back(v);
    adj[v].push_back(u);
  }

  void reset() { fill(visited.begin(), visited.end(), false); }

  ll getAns(int node) { return ans[node]; }

  info dfs(int root) {
    ll sz = 0;
    ll denom = 1;
    ll num = 1;
    visited[root] = true;

    for (auto &i : adj[root]) {
      if (!visited[i]) {
        auto tp_i = dfs(i);
        sz += tp_i.sz;
        num = (num * tp_i.tp) % M;
        denom = (denom * factorial(tp_i.sz)) % M;
      }
    }

    ll tp = (factorial(sz) * num) % M;
    tp = (tp * modinv(denom)) % M;
    sz += 1;
    dp[root] = info{tp, sz};
    return info{tp, sz};
  }

  void calcAns(int root, info dp_parent) {
    visited[root] = true;
    if (root == 1) {
      ans[root] = dp[root].tp;
    }

    for (auto &i : adj[root]) {
      if (!visited[i]) {
        auto dp_p = dp_parent;
        auto dp_c = dp[i];
        // Remove Child
        dp_p.tp = (dp_p.tp * factorial(dp_p.sz - dp_c.sz - 1)) % M;
        dp_p.tp = (dp_p.tp * modinv(factorial(dp_p.sz - 1))) % M;
        dp_p.tp = (dp_p.tp * factorial(dp_c.sz)) % M;
        dp_p.tp = (dp_p.tp * modinv(dp_c.tp)) % M;

        // Calc child
        dp_c.tp = (dp_c.tp * factorial(dp_p.sz - 1)) % M;
        dp_c.tp = (dp_c.tp * modinv(factorial(dp_c.sz - 1))) % M;
        dp_c.tp = (dp_c.tp * dp_p.tp) % M;
        dp_c.tp = (dp_c.tp * modinv(factorial(dp_p.sz - dp_c.sz))) % M;

        dp_c.sz = dp_p.sz;
        ans[i] = dp_c.tp;
        assert(ans[i] >= 0);
        calcAns(i, dp_c);
      }
    }
  }
};

void solve() {
  int n;
  cin >> n;
  Tree t(n);
  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    t.addEdge(u, v);
  }

  auto ans = t.dfs(1);
  t.reset();
  t.calcAns(1, ans);

  for (int i = 1; i <= n; ++i) {
    cout << t.getAns(i) << endl;
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
#ifdef D_LOCAL
  freopen("in.in", "r", stdin);
  freopen("out.out", "w", stdout);
#endif

  solve();

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long int;
constexpr ll M = 1e9 + 7;
inline ll sum_1(const ll n) { return ((n * (n + 1)) / 2) % M; }
inline ll sum_2(const ll n) { return ((n * (n + 1) * (2 * n + 1)) / 6) % M; }
inline ll sum_3(const ll n) {
  return (((n * (n + 1)) / 2) % M * ((n * (n + 1)) / 2) % M) % M;
}
inline ll sum_1(const ll i, const ll j) {
  return (sum_1(j) - sum_1(i - 1) + M) % M;
}
inline ll sum_2(const ll i, const ll j) {
  return (sum_2(j) - sum_2(i - 1) + M) % M;
}
inline ll sum_3(const ll i, const ll j) {
  return (sum_3(j) - sum_3(i - 1) + M) % M;
}

ll pow_f(ll b, ll e) {
  ll ans = 1;
  b = b % M;
  while (e) {
    if (e & 1) {
      ans = (ans * b) % M;
    }
    b = (b * b) % M;
    e >>= 1;
  }
  return ans;
}

ll nCr(ll n, ll r) {
  if (r > n || n < 0 || r < 0) {
    return 0;
  }
  if (n == r || r == 0) {
    return 1;
  }
  r = min(r, n - r);
  ll ans = 1;
  for (ll i = 1; i <= r; ++i) {
    ans = (ans * (n - i + 1)) / (i);
  }
  return ans;
}

ll nCr(tuple<ll, ll> t) {
  auto n = get<0>(t);
  auto k = get<1>(t);
  return nCr(n + k - 1, k);
}

unordered_map<ll, ll> fact;
ll factorial(ll n) {
  if (fact.find(n) != fact.end()) {
    return fact[n];
  }
  ll ans = 1;
  for (ll i = 1; i <= n; ++i) {
    ans = (ans * i) % M;
  }
  fact[n] = ans;
  return ans;
}

inline ll modinv(ll n) { return pow_f(n, M - 2); }

class Tree {
  vector<vector<int>> adj;
  struct info {
    ll tp;
    ll sz;
  };

  vector<bool> visited;

public:
  Tree(int n) : visited(n + 1, false), adj(n + 1) {}

  void addEdge(int u, int v) {
    adj[u].push_back(v);
    adj[v].push_back(u);
  }

  void reset() { fill(visited.begin(), visited.end(), false); }

  info dp(int root) {
    ll sz = 0;
    ll denom = 1;
    ll num = 1;
    visited[root] = true;

    for (auto &i : adj[root]) {
      if (!visited[i]) {
        auto tp_i = dp(i);
        sz += tp_i.sz;
        num = (num * tp_i.tp) % M;
        denom = (denom * factorial(tp_i.sz)) % M;
      }
    }

    ll ans = (factorial(sz) * (num * modinv(denom)) % M) % M;
    sz += 1;
    return info{ans, sz};
  }
};

void solve() {
  int n;
  cin >> n;
  Tree t(n);
  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    t.addEdge(u, v);
  }

  for (int i = 1; i <= n; ++i) {
    auto ans = t.dp(i);
    cout << ans.tp << endl;
    t.reset();
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
#ifdef D_LOCAL
  freopen("in.in", "r", stdin);
  freopen("out.out", "w", stdout);
#endif

  solve();

  return 0;
}