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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

using namespace std;

typedef long long ll;
const ll MOD = (1e9 + 7ll);
const int MAXN = 200001;
ll f[MAXN];
ll finv[MAXN];

ll power(ll x, int p) {
  ll res = 1;
  x = x % MOD;
  if (x == 0)
    return 0;
  while (p > 0) {
    if (p & 1)
      res = (res * x) % MOD;
    p >>= 1;
    x = (x * x) % MOD;
  }
  return res;
}

ll inv(ll x) { return power(x, MOD - 2); }

void table() {
  f[0] = 1;
  finv[0] = 1;
  f[1] = 1;
  finv[1] = 1;
  for (ll i = 2; i < MAXN; i++) {
    f[i] = (i * f[i - 1]) % MOD;
    finv[i] = inv(f[i]);
  }
}

ll mn(int m, int n) { // return (m + n)! / m! n!
  return (((f[m + n] * finv[m]) % MOD) * finv[n]) % MOD;
}

map<pair<int, int>, pair<ll, int>> mem;

pair<ll, int> weight(vector<vector<int>> &tree, int k, int root) {
  ll ans = 1;
  int sz = 0;
  auto it = mem.find({k, root});
  if (it != mem.end())
    return it->second;
  for (auto child : tree[k]) {
    if (child == root)
      continue;
    auto sub = weight(tree, child, k);
    ans = (((ans * sub.first) % MOD) * mn(sz, sub.second)) % MOD;
    sz += sub.second;
  }
  mem[{k, root}] = {ans, sz + 1};
  return {ans, sz + 1};
}

void deduce(vector<vector<int>> &tree, int k, int root, vector<ll> &ans) {
  int n = tree.size();
  auto r = mem[{root, root}];
  auto q = mem[{k, root}];
  // calculate mem[{root, k}] first
  ll root_val = (((((((r.first * inv(q.first)) % MOD) * f[q.second]) % MOD) *
                   f[n - 1 - q.second]) %
                  MOD) *
                 finv[n - 1]) %
                MOD;
  ll root_sz = n - q.second;
  mem[{root, k}] = {root_val, root_sz};

  ll k_val = (((((((q.first * root_val) % MOD) * finv[q.second - 1]) % MOD) *
                f[n - 1]) %
               MOD) *
              finv[root_sz]) %
             MOD;
  mem[{k, k}] = {k_val, n};
  ans[k] = k_val;
  for (auto child : tree[k]) {
    if (child == root)
      continue;
    deduce(tree, child, k, ans);
  }
}

int main() {
  table();
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  int n;
  cin >> n;
  vector<vector<int>> tree(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    tree[a].push_back(b);
    tree[b].push_back(a);
  }

  vector<ll> ans(n);
  auto r = weight(tree, 0, 0);
  ans[0] = r.first;

  for (auto c : tree[0]) {
    deduce(tree, c, 0, ans);
  }
  for (int i = 0; i < n; i++)
    cout << ans[i] << endl;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

using namespace std;

typedef long long ll;
const ll MOD = (1e9 + 7ll);
ll f[200001];
ll finv[200001];

ll power(ll x, int p) {
  ll res = 1;
  x = x % MOD;
  if (x == 0)
    return 0;
  while (p > 0) {
    if (p & 1)
      res = (res * x) % MOD;
    p >>= 1;
    x = (x * x) % MOD;
  }
  return res;
}

ll inv(ll x) { return power(x, MOD - 2); }

void table() {
  ll ans = 1ll;
  f[0] = 1;
  finv[0] = 1;
  f[1] = 1;
  finv[1] = 1;
  for (ll i = 2; i <= 200000; i++) {
    f[i] = (i * f[i - 1]) % MOD;
    finv[i] = inv(f[i]);
  }
}

ll mn(int m, int n) { // return (m + n)! / m! n!
  return (((f[m + n] * finv[m]) % MOD) * finv[n]) % MOD;
}

map<pair<int, int>, pair<ll, int>> mem;

pair<ll, int> weight(vector<vector<int>> &tree, int k, int root) {
  ll ans = 1;
  int sz = 0;
  auto it = mem.find({k, root});
  if (it != mem.end())
    return it->second;
  for (auto child : tree[k]) {
    if (child == root)
      continue;
    auto sub = weight(tree, child, k);
    ans = (((ans * sub.first) % MOD) * mn(sz, sub.second)) % MOD;
    sz += sub.second;
  }
  mem[{k, root}] = {ans, sz + 1};
  return {ans, sz + 1};
}

ll solve(vector<vector<int>> &tree, int k) {
  auto ans = weight(tree, k, -1);
  return ans.first;
}

int main() {
  table();
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  int n;
  cin >> n;
  vector<vector<int>> tree(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    tree[a].push_back(b);
    tree[b].push_back(a);
  }

  for (int k = 0; k < n; k++) {
    cout << solve(tree, k) << endl;
  }
}