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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const ll M = 1000000007;

const int MAX = 200005;
const int MOD = 1000000007;

long long fac[MAX], finv[MAX], inv[MAX];

// テーブルを作る前処理
void COMinit() {
  fac[0] = fac[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < MAX; i++) {
    fac[i] = fac[i - 1] * i % MOD;
    inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    finv[i] = finv[i - 1] * inv[i] % MOD;
  }
}

// 二項係数計算
long long COM(int n, int k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac[n] * (finv[k] * finv[n - k] % MOD) % MOD;
}

template <std::uint_fast64_t Modulus> class modint {
  using u64 = std::uint_fast64_t;

public:
  u64 a;

  constexpr modint(const u64 x = 0) noexcept : a(x % Modulus) {}
  constexpr u64 &value() noexcept { return a; }
  constexpr const u64 &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    u64 exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

using mint = modint<1000000007>;

ll n;
struct edgeinfo {
  ll next_node;
  ll node_count;
  mint result; // dfs2( ?, next_node )
  ll rev;
};
vector<edgeinfo> einfo;
vector<mint> node_result;

vector<vector<ll>> e;

/*
        このノード以下の色の塗り方を返す
*/
mint dfs2(ll node, ll parent, ll node_count, ll i) {

  if (node_result[node].value() && i != -1) {
    edgeinfo &ei = einfo[i ^ 1];
    return node_result[node] / (ei.result * COM(n - 1, ei.node_count));
  }

  mint result = 1;
  ll m = node_count - 1;
  mint rev_result = 0;
  ll rev_node_count = 0;

  for (auto i : e[node]) {
    edgeinfo &ei = einfo[i];
    if (ei.next_node == parent) {
      if (ei.result.value()) {
        rev_result = ei.result;
        rev_node_count = ei.node_count;
      }
      continue;
    }
    if (ei.result.value() == 0) {
      ei.result = dfs2(ei.next_node, node, ei.node_count, i);
    }

    result *= ei.result * COM(m, ei.node_count);
    m -= ei.node_count;
  }

  if (rev_result.value()) {
    node_result[node] = result * rev_result * COM(n - 1, rev_node_count);
  }

  return result;
}

// このノードと子孫ノードの数を数える
ll dfs1(ll node, ll parent) {
  ll result = 1;
  for (auto i : e[node]) {
    edgeinfo &ei = einfo[i];
    edgeinfo &er = einfo[i ^ 1];
    if (ei.next_node == parent) {
      continue;
    }
    ei.node_count = dfs1(ei.next_node, node);
    er.node_count = n - ei.node_count;
    result += ei.node_count;
  }
  return result;
}

int main() {
  COMinit();
  cin >> n;
  e.resize(n);
  node_result.resize(n);

  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    einfo.push_back({b, 0, 0, i * 2 + 1});
    einfo.push_back({a, 0, 0, i * 2});

    e[a].push_back(i * 2);
    e[b].push_back(i * 2 + 1);
  }

  dfs1(0, -1);
  for (int i = 0; i < n; i++) {
    mint result = dfs2(i, -1, n, -1);
    cout << result.value() << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const ll M = 1000000007;

const int MAX = 200005;
const int MOD = 1000000007;

long long fac[MAX], finv[MAX], inv[MAX];

// テーブルを作る前処理
void COMinit() {
  fac[0] = fac[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < MAX; i++) {
    fac[i] = fac[i - 1] * i % MOD;
    inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    finv[i] = finv[i - 1] * inv[i] % MOD;
  }
}

// 二項係数計算
long long COM(int n, int k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac[n] * (finv[k] * finv[n - k] % MOD) % MOD;
}

template <std::uint_fast64_t Modulus> class modint {
  using u64 = std::uint_fast64_t;

public:
  u64 a;

  constexpr modint(const u64 x = 0) noexcept : a(x % Modulus) {}
  constexpr u64 &value() noexcept { return a; }
  constexpr const u64 &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    u64 exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

using mint = modint<1000000007>;

ll n;
struct edgeinfo {
  ll next_node;
  ll node_count;
  mint result; // dfs2( ?, next_node )
  ll rev;
};
vector<edgeinfo> einfo;
vector<mint> node_result;

vector<vector<ll>> e;

/*
        このノード以下の色の塗り方を返す
*/
mint dfs2(ll node, ll parent, ll node_count, ll i) {

  if (node_result[node].value() && i != -1) {
    edgeinfo &ei = einfo[i ^ 1];
    return node_result[node] / (ei.result * COM(n - 1, ei.node_count));
  }

  mint result = 1;
  ll m = node_count - 1;
  mint rev_result = 0;
  ll rev_node_count = 0;

  for (auto i : e[node]) {
    edgeinfo &ei = einfo[i];
    if (ei.next_node == parent) {
      if (ei.result.value()) {
        rev_result = ei.result;
        rev_node_count = ei.node_count;
      }
      continue;
    }
    if (ei.result.value() == 0) {
      ei.result = dfs2(ei.next_node, node, ei.node_count, i);
    }

    result *= ei.result * COM(m, ei.node_count);
    m -= ei.node_count;
  }

  if (rev_result.value()) {
    node_result[node] = result * rev_result * COM(n - 1, rev_node_count);
  }

  return result;
}

// このノードと子孫ノードの数を数える
ll dfs1(ll node, ll parent) {
  ll result = 1;
  for (auto i : e[node]) {
    edgeinfo &ei = einfo[i];
    if (ei.next_node == parent)
      continue;
    if (ei.node_count == 0)
      ei.node_count = dfs1(ei.next_node, node);
    result += ei.node_count;
  }
  return result;
}

int main() {
  COMinit();
  cin >> n;
  e.resize(n);
  node_result.resize(n);

  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    einfo.push_back({b, 0, 0, i * 2 + 1});
    einfo.push_back({a, 0, 0, i * 2});

    e[a].push_back(i * 2);
    e[b].push_back(i * 2 + 1);
  }

  for (int i = 0; i < n; i++) {
    dfs1(i, -1);
  }

  for (int i = 0; i < n; i++) {
    mint result = dfs2(i, -1, n, -1);
    cout << result.value() << endl;
  }

  return 0;
}