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

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

constexpr int64_t MOD = 1e9 + 7;
constexpr int64_t MAX_N = 2e5;

// a^b % modを求める
int64_t MODpow(int64_t a, int64_t b, int64_t mod) {
  int64_t result = 1;
  while (b) {
    if (b % 2 == 1) {
      result *= a;
      result %= mod;
    }

    b /= 2;
    a *= a;
    a %= mod;
  }

  return result;
}

class Combination {
public:
  Combination(int64_t max_num, int64_t mod) {
    fact.resize(max_num + 1, 1);
    inv_fact.resize(max_num + 1, 1);
    mod_ = mod;
    for (int64_t i = 2; i <= max_num; i++) {
      fact[i] = i * fact[i - 1] % mod_;
      inv_fact[i] = MODpow(fact[i], mod_ - 2, mod_);
      assert(fact[i] * inv_fact[i] % mod_ == 1);
    }
  }
  int64_t operator()(int64_t n, int64_t m) const {
    if (m < 0 || m > n)
      return 0;
    return fact[n] * inv_fact[n - m] % mod_ * inv_fact[m] % mod_;
  }
  vector<int64_t> fact, inv_fact;

private:
  int64_t mod_;
};
Combination comb(MAX_N, MOD);

struct DP_Element {
  int64_t dp, size;
  DP_Element() : dp(1), size(1) {}
  DP_Element operator+=(const DP_Element &rhs) {
    (dp *= rhs.dp) %= MOD;
    (dp *= comb.inv_fact[rhs.size]) %= MOD;
    size += rhs.size;
    return *this;
  }
};

int main() {
  int64_t N;
  cin >> N;
  vector<vector<int64_t>> connect(N);
  for (int64_t i = 0; i < N - 1; i++) {
    int64_t a, b;
    cin >> a >> b;
    a--;
    b--;
    connect[a].push_back(b);
    connect[b].push_back(a);
  }

  // dp1[i] := 0を根としたときのノードi以下の部分木についての答え
  vector<DP_Element> dp1(N);
  function<DP_Element(int64_t, int64_t)> f1 = [&](int64_t curr, int64_t pre) {
    for (int64_t next_node : connect[curr]) {
      if (next_node == pre) {
        continue;
      }

      DP_Element next_dp = f1(next_node, curr);
      dp1[curr] += next_dp;
    }

    (dp1[curr].dp *= comb.fact[dp1[curr].size - 1]) %= MOD;
    return dp1[curr];
  };

  f1(0, -1);

  // dp2[i] := iを根としたときのiについての答え
  vector<DP_Element> dp2(N);
  dp2[0] = dp1[0];
  function<void(int64_t, int64_t)> f2 = [&](int64_t curr, int64_t pre) {
    for (int64_t next_node : connect[curr]) {
      if (next_node == pre) {
        continue;
      }

      DP_Element add = dp2[curr];
      // dp1[next_node]からの影響を取り除く
      (add.dp *= comb.inv_fact[add.size - 1]) %= MOD;
      (add.dp *= MODpow(dp1[next_node].dp, MOD - 2, MOD)) %= MOD;
      (add.dp *= comb.fact[dp1[next_node].size]) %= MOD;
      add.size -= dp1[next_node].size;
      (add.dp *= comb.fact[add.size - 1]) %= MOD;

      // dp2[next_node]に足す
      DP_Element &target = dp2[next_node];
      target = dp1[next_node];
      (target.dp *= comb.inv_fact[target.size - 1]) %= MOD;
      target += add;
      (target.dp *= comb.fact[target.size - 1]) %= MOD;

      // 次へ遷移
      f2(next_node, curr);
    }
  };

  f2(0, -1);

  for (int64_t i = 0; i < N; i++) {
    cout << dp2[i].dp << endl;
  }
}

Error: TLE

Fixed submission:

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

int64_t N;
vector<vector<int64_t>> connect;

constexpr int64_t MOD = 1e9 + 7;
class Combination {
public:
  Combination(int64_t max_num, int64_t mod) {
    fact_.resize(max_num + 1, 1);
    inv_fact_.resize(max_num + 1, 1);
    mod_ = mod;
    for (int64_t i = 2; i <= max_num; i++) {
      fact_[i] = i * fact_[i - 1] % mod_;
      inv_fact_[i] = MODpow(fact_[i], mod_ - 2);
      assert(fact_[i] * inv_fact_[i] % mod_ == 1);
    }
  }
  int64_t operator()(int64_t n, int64_t m) const {
    if (m < 0 || m > n)
      return 0;
    return fact_[n] * inv_fact_[n - m] % mod_ * inv_fact_[m] % mod_;
  }

private:
  int64_t MODpow(int64_t n, int64_t m) const {
    int64_t result = 1;
    while (m) {
      if (m % 2 == 1) {
        result *= n;
        result %= mod_;
      }

      m /= 2;
      n *= n;
      n %= mod_;
    }

    return result;
  }

  vector<int64_t> fact_, inv_fact_;
  int64_t mod_;
};
Combination comb(200000, MOD);

vector<vector<int64_t>> memo, child_num;

int64_t countChildNum(int64_t node, int64_t pre_node, int64_t index) {
  if (pre_node != -1 && child_num[pre_node][index] != -1) {
    return child_num[pre_node][index];
  }

  int64_t result = 0;
  for (int64_t i = 0; i < connect[node].size(); i++) {
    if (connect[node][i] == pre_node) {
      continue;
    }
    result += countChildNum(connect[node][i], node, i) + 1;
  }
  if (pre_node != -1) {
    child_num[pre_node][index] = result;
  }
  return result;
}

int64_t solve(int64_t node, int64_t pre_node, int64_t index) {
  if (pre_node != -1 && memo[pre_node][index] != -1) {
    return memo[pre_node][index];
  }
  int64_t result = 1;

  int64_t all_child_num = countChildNum(node, pre_node, index);

  for (int64_t i = 0; i < connect[node].size(); i++) {
    int64_t next_node = connect[node][i];
    if (next_node == pre_node) {
      continue;
    }

    int64_t curr_child_num = countChildNum(next_node, node, i) + 1;
    int64_t curr_value = solve(next_node, node, i);
    (result *= comb(all_child_num, curr_child_num) * curr_value % MOD) %= MOD;
    all_child_num -= curr_child_num;
  }

  if (pre_node != -1) {
    memo[pre_node][index] = result;
  }
  return result;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cin >> N;
  connect.resize(N);
  memo.resize(N);
  child_num.resize(N);
  for (int64_t i = 0; i < N - 1; i++) {
    int64_t a, b;
    cin >> a >> b;
    a--;
    b--;
    connect[a].push_back(b);
    memo[a].push_back(-1);
    child_num[a].push_back(-1);

    connect[b].push_back(a);
    memo[b].push_back(-1);
    child_num[b].push_back(-1);
  }

  for (int64_t i = 0; i < N; i++) {
    cout << solve(i, -1, -1) << endl;
  }
}