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

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

#define rep(i, x, y) for (ll i = (x); i < (y); i++)
#define rrep(i, x, y) for (ll i = (ll)(y)-1; i >= (x); i--)
#define all(x) (x).begin(), (x).end()
#ifdef LOCAL
#define dump(x) cerr << #x << " = " << (x) << endl
#define debug(x)                                                               \\
  cerr << #x << " = " << (x) << " (L" << __LINE__ << ")"                       \\
       << " " << __FILE__ << endl
#define dassert(...) assert(__VA_ARGS__)
#else
#define dump(x)
#define debug(x)
#define dassert(...)
#endif
// #define int long long

typedef long long ll;
const ll MOD = 1e9 + 7;
const long double EPS = 1e-8;

long long modPow(long long a, long long x, long long p) {
  long long res = 1;
  while (x > 0) {
    if (x % 2 != 0) {
      res = (res * a) % p;
    }
    a = (a * a) % p;
    x /= 2;
  }
  return res;
}

long long modInverse(long long a, long long p) {
  // calculates the modular multiplicative of a mod m.
  //(assuming p is prime).
  return modPow(a, p - 2, p);
}

long long modFactorial(long long a, long long p) {
  long long fact = 1;
  for (long long i = 2; i <= a; i++) {
    fact *= i;
    fact %= p;
  }
  return fact;
}

struct combination {
  void preprocess(long long max_n, long long p) {
    if (_is_preprocessed)
      return;

    factorials_ = vector<ll>(max_n + 1, 1);
    inverse_factorials_ = vector<ll>(max_n + 1, 1);
    for (long long i = 1; i < max_n + 1; i++) {
      factorials_[i] = factorials_[i - 1] * i;
      factorials_[i] %= p;
    }
    inverse_factorials_[max_n] = modPow(factorials_[max_n], p - 2, p);
    for (long long i = max_n - 1; i >= 1; i--) {
      inverse_factorials_[i] = inverse_factorials_[i + 1] * (i + 1);
      inverse_factorials_[i] %= MOD;
    }

    _is_preprocessed = true;
  }

  long long combination_with_preprocessing(long long n, long long k,
                                           long long p) {
    if (!_is_preprocessed)
      assert(false);

    long long result = factorials_[n];
    result *= inverse_factorials_[k];
    result %= p;
    result *= inverse_factorials_[n - k];
    result %= p;
    return result;
  }

  static long long combination_without_preprocessing(long long n, long long k,
                                                     long long p) {
    // calculates C(n,k) mod p (assuming p is prime).
    long long numerator = 1; // n * (n-1) * ... * (n-k+1)
    for (int i = 0; i < k; i++) {
      numerator = (numerator * (n - i)) % p;
    }

    long long denominator = 1; // k!
    for (int i = 1; i <= k; i++) {
      denominator = (denominator * i) % p;
    }
    // numerator / denominator mod p.
    return (numerator * modInverse(denominator, p)) % p;
  }

private:
  vector<ll> factorials_, inverse_factorials_;
  bool _is_preprocessed = false;
};

pair<ll, ll> dfs(int root, int current, vector<vector<int>> &edges,
                 vector<unordered_map<int, pair<ll, ll>>> &dp,
                 vector<pair<ll, ll>> &vdp) {
  static combination comb;
  comb.preprocess(edges.size() + 1, MOD);

  if (root != -1 && dp[root].find(current) != dp[root].end()) {
    return dp[root][current];
  } else if (vdp[current].first != 0) {
    ll numVertex = vdp[current].first, numCase = vdp[current].second;
    auto diff = dfs(current, root, edges, dp, vdp);
    numCase *= modInverse(
        comb.combination_with_preprocessing(numVertex, diff.first, MOD), MOD);
    numCase %= MOD;
    numCase *= modInverse(diff.second, MOD);
    numCase %= MOD;
    return dp[root][current] = {numVertex - diff.first + 1, numCase};
  }

  vector<pair<ll, ll>> values;
  for (const auto &next : edges[current]) {
    if (next == root)
      continue;
    if (dp[current].find(next) != dp[current].end()) {
      values.emplace_back(dp[current][next]);
    } else {
      values.emplace_back(dfs(current, next, edges, dp, vdp));
    }
  }
  ll numCase = 1, numVertex = 0;
  for (const auto &value : values)
    numVertex += value.first;
  ll nv = numVertex;
  for (const auto &value : values) {
    numCase *= value.second;
    numCase %= MOD;
    numCase *= comb.combination_with_preprocessing(nv, value.first, MOD);
    numCase %= MOD;
    nv -= value.first;
  }
  if (root != -1) {
    return dp[root][current] = {numVertex + 1, numCase};
  } else {
    return vdp[current] = {numVertex, numCase};
  }
}

void solve(long long N, std::vector<long long> a, std::vector<long long> b) {
  vector<unordered_map<int, pair<ll, ll>>> dp(
      N + 1);                      // dp[from][to] = {vertex cnt, case cnt};
  vector<pair<ll, ll>> vdp(N + 1); // dp[vertex index] = {vertex cnt, case cnt};
  vector<vector<int>> edges(N + 1);
  rep(i, 0, N - 1) {
    edges[a[i]].emplace_back(b[i]);
    edges[b[i]].emplace_back(a[i]);
  }
  rep(i, 1, N + 1) {
    if (edges[i].size() != 1)
      continue;
    dp[edges[i][0]][i] = {1, 1};
  }
  vector<int> dfsOrder(N);
  rep(i, 0, N) dfsOrder[i] = i + 1;
  sort(all(dfsOrder), [&](const int a, const int b) {
    return edges[a].size() > edges[b].size();
  });
  rep(i, 0, N) { dfs(-1, dfsOrder[i], edges, dp, vdp); }
  rep(k, 1, N + 1) { cout << vdp[k].second << endl; }
}

signed main() {
  // ios_base::sync_with_stdio(false);
  // cin.tie(NULL);
  long long N;
  scanf("%lld", &N);
  std::vector<long long> a(N - 1);
  std::vector<long long> b(N - 1);
  for (int i = 0; i < N - 1; i++) {
    scanf("%lld", &a[i]);
    scanf("%lld", &b[i]);
  }
  solve(N, std::move(a), std::move(b));
  return 0;
}


Error: TLE

Fixed submission:

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

#define rep(i, x, y) for (ll i = (x); i < (y); i++)
#define rrep(i, x, y) for (ll i = (ll)(y)-1; i >= (x); i--)
#define all(x) (x).begin(), (x).end()
#ifdef LOCAL
#define dump(x) cerr << #x << " = " << (x) << endl
#define debug(x)                                                               \\
  cerr << #x << " = " << (x) << " (L" << __LINE__ << ")"                       \\
       << " " << __FILE__ << endl
#define dassert(...) assert(__VA_ARGS__)
#else
#define dump(x)
#define debug(x)
#define dassert(...)
#endif
// #define int long long

typedef long long ll;
const ll MOD = 1e9 + 7;
const long double EPS = 1e-8;

long long modPow(long long a, long long x, long long p) {
  long long res = 1;
  while (x > 0) {
    if (x % 2 != 0) {
      res = (res * a) % p;
    }
    a = (a * a) % p;
    x /= 2;
  }
  return res;
}

long long modInverse(long long a, long long p) {
  // calculates the modular multiplicative of a mod m.
  //(assuming p is prime).
  return modPow(a, p - 2, p);
}

long long modFactorial(long long a, long long p) {
  long long fact = 1;
  for (long long i = 2; i <= a; i++) {
    fact *= i;
    fact %= p;
  }
  return fact;
}

struct combination {
  void preprocess(long long max_n, long long p) {
    if (_is_preprocessed)
      return;

    factorials_ = vector<ll>(max_n + 1, 1);
    inverse_factorials_ = vector<ll>(max_n + 1, 1);
    for (long long i = 1; i < max_n + 1; i++) {
      factorials_[i] = factorials_[i - 1] * i;
      factorials_[i] %= p;
    }
    inverse_factorials_[max_n] = modPow(factorials_[max_n], p - 2, p);
    for (long long i = max_n - 1; i >= 1; i--) {
      inverse_factorials_[i] = inverse_factorials_[i + 1] * (i + 1);
      inverse_factorials_[i] %= MOD;
    }

    _is_preprocessed = true;
  }

  long long combination_with_preprocessing(long long n, long long k,
                                           long long p) {
    if (!_is_preprocessed)
      assert(false);

    long long result = factorials_[n];
    result *= inverse_factorials_[k];
    result %= p;
    result *= inverse_factorials_[n - k];
    result %= p;
    return result;
  }

  static long long combination_without_preprocessing(long long n, long long k,
                                                     long long p) {
    // calculates C(n,k) mod p (assuming p is prime).
    long long numerator = 1; // n * (n-1) * ... * (n-k+1)
    for (int i = 0; i < k; i++) {
      numerator = (numerator * (n - i)) % p;
    }

    long long denominator = 1; // k!
    for (int i = 1; i <= k; i++) {
      denominator = (denominator * i) % p;
    }
    // numerator / denominator mod p.
    return (numerator * modInverse(denominator, p)) % p;
  }

private:
  vector<ll> factorials_, inverse_factorials_;
  bool _is_preprocessed = false;
};

pair<ll, ll> dfs(int root, int current, vector<vector<int>> &edges,
                 vector<unordered_map<int, pair<ll, ll>>> &dp,
                 vector<pair<ll, ll>> &vdp) {
  static combination comb;
  comb.preprocess(edges.size() + 1, MOD);

  if (root != -1 && dp[root].find(current) != dp[root].end()) {
    return dp[root][current];
  } else if (vdp[current].first != 0) {
    ll numVertex = vdp[current].first, numCase = vdp[current].second;
    auto diff = dfs(current, root, edges, dp, vdp);
    numCase *= modInverse(
        comb.combination_with_preprocessing(numVertex, diff.first, MOD), MOD);
    numCase %= MOD;
    numCase *= modInverse(diff.second, MOD);
    numCase %= MOD;
    return dp[root][current] = {numVertex - diff.first + 1, numCase};
  }

  vector<pair<ll, ll>> values;
  for (const auto &next : edges[current]) {
    if (next == root)
      continue;
    if (dp[current].find(next) != dp[current].end()) {
      values.emplace_back(dp[current][next]);
    } else {
      values.emplace_back(dfs(current, next, edges, dp, vdp));
    }
  }
  ll numCase = 1, numVertex = 0;
  for (const auto &value : values)
    numVertex += value.first;
  ll nv = numVertex;
  for (const auto &value : values) {
    numCase *= value.second;
    numCase %= MOD;
    numCase *= comb.combination_with_preprocessing(nv, value.first, MOD);
    numCase %= MOD;
    nv -= value.first;
  }
  if (root != -1) {
    return dp[root][current] = {numVertex + 1, numCase};
  } else {
    return vdp[current] = {numVertex, numCase};
  }
}

void solve(long long N, std::vector<long long> a, std::vector<long long> b) {
  vector<unordered_map<int, pair<ll, ll>>> dp(
      N + 1);                      // dp[from][to] = {vertex cnt, case cnt};
  vector<pair<ll, ll>> vdp(N + 1); // dp[vertex index] = {vertex cnt, case cnt};
  vector<vector<int>> edges(N + 1);
  rep(i, 0, N - 1) {
    edges[a[i]].emplace_back(b[i]);
    edges[b[i]].emplace_back(a[i]);
  }
  rep(i, 1, N + 1) {
    if (edges[i].size() != 1)
      continue;
    dp[edges[i][0]][i] = {1, 1};
  }
  rep(k, 1, N + 1) { cout << dfs(-1, k, edges, dp, vdp).second << endl; }
}

signed main() {
  // ios_base::sync_with_stdio(false);
  // cin.tie(NULL);
  long long N;
  scanf("%lld", &N);
  std::vector<long long> a(N - 1);
  std::vector<long long> b(N - 1);
  for (int i = 0; i < N - 1; i++) {
    scanf("%lld", &a[i]);
    scanf("%lld", &b[i]);
  }
  solve(N, std::move(a), std::move(b));
  return 0;
}