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

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (ll i = 0; i < (n); ++i)

struct mint {
  static const long long MOD = 1000000007;

  long long v;

  mint() : v(0) {}
  mint(long long x) {
    x = x % MOD;
    if (x < 0) {
      x += MOD;
    }
    v = x;
  }
  mint &operator+=(mint a) {
    v += a.v;
    if (v >= MOD) {
      v -= MOD;
    }
    return *this;
  }
  mint &operator*=(mint a) {
    v *= a.v;
    v %= MOD;
    return *this;
  }
  mint &operator-=(mint a) {
    v -= a.v;
    if (v < 0) {
      v += MOD;
    }
    return *this;
  }
  mint &operator/=(mint a) { return (*this) *= a.inv(); }

  mint operator*(mint a) const { return mint(v) *= a; }

  mint operator+(mint a) const { return mint(v) += a; }

  mint operator-(mint a) const { return mint(v) -= a; }

  mint operator/(mint a) const { return mint(v) /= a; }

  mint pow(long long k) const {
    mint res(1), tmp(v);
    while (k) {
      if (k & 1)
        res *= tmp;
      tmp *= tmp;
      k >>= 1;
    }
    return res;
  }
  mint inv() { return pow(MOD - 2); }

  static mint comb(long long n, int k) {
    mint res(1);
    for (int i = 0; i < k; ++i) {
      res *= mint(n - i);
      res /= mint(i + 1);
    }
    return res;
  }

  static mint factorial(long long n) {
    mint res(1);
    for (int i = n; i > 1; --i) {
      res *= mint(i);
    }
    return res;
  }

  bool operator<(const mint &a) const { return v < a.v; };
};

mint operator*(long long l, mint r) { return r * l; }

std::ostream &operator<<(std::ostream &stream, const mint &m) {
  stream << m.v;
  return stream;
}

struct mintcomb {
  vector<mint> fact;
  vector<mint> ifact;
  mintcomb(ll n) : fact(n + 1, 1), ifact(n + 1, 1) {
    for (ll i = 2; i <= n; ++i) {
      fact.at(i) = fact.at(i - 1) * i;
      ifact.at(i) = fact.at(i).inv();
    }
  }
  mint factorial(ll k) const { return fact.at(k); }
  mint combination(ll k, ll l) const {
    return fact.at(k) * ifact.at(l) * ifact.at(k - l);
  }
};

struct HashPair {
  template <class T1, class T2> size_t operator()(const pair<T1, T2> &p) const {
    auto hash1 = hash<T1>{}(p.first);
    auto hash2 = hash<T2>{}(p.second);
    size_t seed = 0;
    seed ^= hash1 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    seed ^= hash2 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    return seed;
  }
};

ll n;
vector<vector<ll>> net;

unordered_map<pair<ll, ll>, ll, HashPair> subgraph_size;

ll init_subgraph_size(ll cur, ll par) {
  ll ret = 1;
  for (ll nxt : net.at(cur)) {
    if (nxt == par) {
      continue;
    }
    ret += init_subgraph_size(nxt, cur);
  }
  if (par != -1) {
    subgraph_size[make_pair(cur, par)] = ret;
    subgraph_size[make_pair(par, cur)] = n - ret;
  }
  return ret;
}

vector<mint> iprodfact;

void init_iprodfact(const mintcomb &mc) {
  iprodfact.resize(n);
  REP(i, n) {
    mint tmp = 1;
    for (ll nxt : net.at(i)) {
      tmp *= mc.factorial(subgraph_size.at(make_pair(i, nxt)));
    }
    iprodfact.at(i) = tmp.inv();
  }
}

unordered_map<pair<ll, ll>, mint, HashPair> cache;
vector<ll> visit_count;
vector<mint> all_prod;

mint solve(ll cur, ll par, const mintcomb &mc) {
  if (cache.count(make_pair(cur, par))) {
    return cache.at(make_pair(cur, par));
  }

  visit_count.at(cur)++;

  mint nump = 1;
  if (visit_count.at(cur) <= 2) {
    for (ll nxt : net.at(cur)) {
      if (nxt == par) {
        continue;
      }
      nump *= solve(nxt, cur, mc);
    }
  } else {
    if (par != -1) {
      nump = all_prod.at(cur) / solve(par, cur, mc);
    } else {
      nump = all_prod.at(cur);
    }
  }

  if (visit_count.at(cur) == 2) {
    if (par != -1) {
      all_prod.at(cur) = nump * solve(par, cur, mc);
    } else {
      all_prod.at(cur) = nump;
    }
  }

  ll n0 = 0;
  if (par != -1) {
    n0 = subgraph_size.at(make_pair(par, cur));
  }
  mint tmp = mc.factorial(n0) * mc.factorial(n - 1 - n0) * iprodfact.at(cur);
  nump *= tmp;

  cache[make_pair(cur, par)] = nump;
  return cache.at(make_pair(cur, par));
}

int main() {
  cin >> n;
  net.resize(n);
  REP(i, n - 1) {
    ll ai, bi;
    cin >> ai >> bi;
    ai--;
    bi--;
    net.at(ai).push_back(bi);
    net.at(bi).push_back(ai);
  }

  mintcomb mc(n);
  init_subgraph_size(0, -1);
  init_iprodfact(mc);

  visit_count.resize(n);
  fill(visit_count.begin(), visit_count.end(), 0);
  all_prod.resize(n);

  REP(i, n) { cout << solve(i, -1, mc) << endl; }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (ll i = 0; i < (n); ++i)

struct mint {
  static const long long MOD = 1000000007;

  long long v;

  mint() : v(0) {}
  mint(long long x) {
    x = x % MOD;
    if (x < 0) {
      x += MOD;
    }
    v = x;
  }
  mint &operator+=(mint a) {
    v += a.v;
    if (v >= MOD) {
      v -= MOD;
    }
    return *this;
  }
  mint &operator*=(mint a) {
    v *= a.v;
    v %= MOD;
    return *this;
  }
  mint &operator-=(mint a) {
    v -= a.v;
    if (v < 0) {
      v += MOD;
    }
    return *this;
  }
  mint &operator/=(mint a) { return (*this) *= a.inv(); }

  mint operator*(mint a) const { return mint(v) *= a; }

  mint operator+(mint a) const { return mint(v) += a; }

  mint operator-(mint a) const { return mint(v) -= a; }

  mint operator/(mint a) const { return mint(v) /= a; }

  mint pow(long long k) const {
    mint res(1), tmp(v);
    while (k) {
      if (k & 1)
        res *= tmp;
      tmp *= tmp;
      k >>= 1;
    }
    return res;
  }
  mint inv() { return pow(MOD - 2); }

  static mint comb(long long n, int k) {
    mint res(1);
    for (int i = 0; i < k; ++i) {
      res *= mint(n - i);
      res /= mint(i + 1);
    }
    return res;
  }

  static mint factorial(long long n) {
    mint res(1);
    for (int i = n; i > 1; --i) {
      res *= mint(i);
    }
    return res;
  }

  bool operator<(const mint &a) const { return v < a.v; };
};

mint operator*(long long l, mint r) { return r * l; }

std::ostream &operator<<(std::ostream &stream, const mint &m) {
  stream << m.v;
  return stream;
}

struct mintcomb {
  vector<mint> fact;
  vector<mint> ifact;
  mintcomb(ll n) : fact(n + 1, 1), ifact(n + 1, 1) {
    for (ll i = 2; i <= n; ++i) {
      fact.at(i) = fact.at(i - 1) * i;
      ifact.at(i) = fact.at(i).inv();
    }
  }
  mint factorial(ll k) const { return fact.at(k); }
  mint combination(ll k, ll l) const {
    return fact.at(k) * ifact.at(l) * ifact.at(k - l);
  }
};

struct HashPair {
  template <class T1, class T2> size_t operator()(const pair<T1, T2> &p) const {
    auto hash1 = hash<T1>{}(p.first);
    auto hash2 = hash<T2>{}(p.second);
    size_t seed = 0;
    seed ^= hash1 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    seed ^= hash2 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    return seed;
  }
};

ll n;
vector<vector<ll>> net;

unordered_map<pair<ll, ll>, ll, HashPair> count_cache;

ll solve_count(ll cur, ll par) {
  if (count_cache.count(make_pair(cur, par))) {
    return count_cache.at(make_pair(cur, par));
  }
  ll ret = 1;
  for (ll nxt : net.at(cur)) {
    if (nxt == par) {
      continue;
    }
    ret += solve_count(nxt, cur);
  }
  count_cache[make_pair(cur, par)] = ret;
  count_cache[make_pair(par, cur)] = n - ret;
  return ret;
}

vector<mint> prodfact, iprodfact;

void calc_prodfact(const mintcomb &mc) {
  prodfact.resize(n);
  iprodfact.resize(n);
  REP(i, n) {
    prodfact.at(i) = 1;
    for (ll nxt : net.at(i)) {
      ll tmp = count_cache.at(make_pair(i, nxt));
      prodfact.at(i) *= mc.factorial(tmp);
    }
    iprodfact.at(i) = prodfact.at(i).inv();
  }
}

unordered_map<pair<ll, ll>, mint, HashPair> cache;

mint solve(ll cur, ll par, const mintcomb &mc) {
  if (cache.count(make_pair(cur, par))) {
    return cache.at(make_pair(cur, par));
  }

  mint nump = 1;
  for (ll nxt : net.at(cur)) {
    if (nxt == par) {
      continue;
    }
    nump *= solve(nxt, cur, mc);
  }

  ll n0 = 0;
  if (par != -1) {
    n0 = count_cache.at(make_pair(par, cur));
  }
  mint tmp = mc.factorial(n0) * mc.factorial(n - 1 - n0) * iprodfact.at(cur);
  nump *= tmp;

  cache[make_pair(cur, par)] = nump;
  return cache.at(make_pair(cur, par));
}

int main() {
  cin >> n;
  net.resize(n);
  REP(i, n - 1) {
    ll ai, bi;
    cin >> ai >> bi;
    ai--;
    bi--;
    net.at(ai).push_back(bi);
    net.at(bi).push_back(ai);
  }

  mintcomb mc(n);
  solve_count(0, -1);
  calc_prodfact(mc);

  REP(i, n) {
    // cout << solve(i, -1, mc).first << endl;
    cout << solve(i, -1, mc) << endl;
  }
  return 0;
}