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

#include <bits/stdc++.h>

using namespace std;

long long prime = 1000000007;

long long pow_mod(long long a, long long r) {
  if (r == 0)
    return 1LL;
  if (r == 1LL)
    return a;
  long long y = pow_mod(a, r / 2);
  long long z = y * y % prime;
  if (r % 2 == 0) {
    return z;
  }
  return z * a % prime;
}

long long inv(long long x) { return pow_mod(x, prime - 2); }

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n;
  cin >> n;
  vector<vector<int>> adj_list(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    adj_list[a].push_back(b);
    adj_list[b].push_back(a);
  }
  vector<long long> fac(n + 1);
  fac[0] = 1;
  // inv_fac[0] = 1;
  for (int i = 1; i <= n; ++i) {
    fac[i] = i * fac[i - 1] % prime;
    // inv_fac[i] = inv(fac[i]);
  }

  vector<int> subtree_sizes(n), xs(n);

  function<void(int, int)> down = [&](int i, int parent) {
    long long x = 1LL;
    int s = 0;
    for (int j : adj_list[i]) {
      if (j != parent) {
        down(j, i);
        x = x * xs[j] % prime;
        s += subtree_sizes[j];
      }
    }
    x = x * (s + 1) % prime;
    xs[i] = (int)x;
    subtree_sizes[i] = s + 1;
  };
  down(0, -1);

  function<void(int, int)> up = [&](int i, int parent) {
    if (parent != -1) {
      xs[i] = xs[parent] * inv(subtree_sizes[i]) % prime *
              (n - subtree_sizes[i]) % prime;
    }
    for (int j : adj_list[i]) {
      if (j != parent) {
        up(j, i);
      }
    }
  };
  up(0, -1);

  for (int i = 0; i < n; ++i) {
    cout << (fac[n] * inv(xs[i]) % prime) << "\\n";
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

long long prime = 1000000007;

long long pow_mod(long long a, long long r) {
  if (r == 0)
    return 1LL;
  if (r == 1LL)
    return a;
  long long y = pow_mod(a, r / 2);
  long long z = y * y % prime;
  if (r % 2 == 0) {
    return z;
  }
  return z * a % prime;
}

long long inv(long long x) { return pow_mod(x, prime - 2); }

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n;
  cin >> n;
  vector<vector<int>> adj_list(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    adj_list[a].push_back(b);
    adj_list[b].push_back(a);
  }
  vector<long long> fac(n + 1), inv_fac(n + 1);
  fac[0] = 1;
  inv_fac[0] = 1;
  for (int i = 1; i <= n; ++i) {
    fac[i] = i * fac[i - 1] % prime;
    inv_fac[i] = inv(fac[i]);
  }

  map<pair<int, int>, pair<int, int>> size_and_choices;
  function<pair<int, int>(int, int)> get_size_and_choices = [&](int i,
                                                                int parent) {
    auto p = make_pair(i, parent);
    if (size_and_choices.find(p) != size_and_choices.end()) {
      return size_and_choices[p];
    }
    long long choices = 1LL;
    int s = 0;
    for (int j : adj_list[i]) {
      if (j != parent) {
        auto rec = get_size_and_choices(j, i);
        choices = choices * ((long long)rec.second) % prime;
        choices = choices * inv_fac[rec.first] % prime;
        s += rec.first;
      }
    }
    choices = choices * fac[s] % prime;
    return size_and_choices[p] = make_pair(s + 1, (int)choices);
  };

  for (int i = 0; i < n; ++i) {
    cout << get_size_and_choices(i, -1).second << "\\n";
  }

  return 0;
}