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

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

typedef long long ll;
typedef vector<ll> vll;
typedef vector<vector<ll>> vvll;
typedef vector<vector<vector<ll>>> vvvll;
typedef vector<bool> vb;
typedef vector<vector<bool>> vvb;
typedef vector<vector<vector<bool>>> vvvb;
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef vector<vpll> vvpll;
typedef vector<double> vd;
typedef vector<vd> vdd;
typedef pair<ll, vll> plvll;

#define FOR(i, x, y) for (ll i = (ll)x; i < (ll)y; ++i)
#define REP(i, y) FOR(i, 0, y)
#define RFOR(i, x, y) for (ll i = (ll)x; i >= (ll)y; --i)
#define RREP(i, x) RFOR(i, x, 0)
#define ALL(a) a.begin(), a.end()
#define pb push_back
#define debug_print(x...)                                                      \\
  cerr << "line " << __LINE__ << " : ";                                        \\
  debug_print_in(x);

template <typename First> void debug_print_in(First first) {
  cerr << first << endl;
  return;
}

template <typename First, typename... Rest>
void debug_print_in(First first, Rest... rest) {
  cerr << first << " ";
  debug_print_in(rest...);
  return;
}

void IN(void) { return; }

template <typename First, typename... Rest>
void IN(First &first, Rest &...rest) {
  cin >> first;
  IN(rest...);
  return;
}

template <typename First> void OUT(First first) {
  cout << first << endl;
  return;
}

template <typename First, typename... Rest>
void OUT(First first, Rest... rest) {
  cout << first << " ";
  OUT(rest...);
  return;
}

template <class t, class u> void chmax(t &a, u b) {
  if (a < b)
    a = b;
};
template <class t, class u> void chmin(t &a, u b) {
  if (a > b)
    a = b;
};
int popcount(int t) { return __builtin_popcount(t); }  // GCC
int popcount(ll t) { return __builtin_popcountll(t); } // GCC

template <typename T> void vec_print(vector<T> VEC) {
  REP(i, VEC.size()) { cerr << VEC[i] << " "; }
  cerr << endl;
};

template <typename T>
void mat_print(vector<vector<T>> MAT){
    REP(i, MAT.size()){REP(j, MAT[i].size()){cerr << MAT[i][j] << " ";
}
cerr << endl;
}
}
;

constexpr int INF = (1 << 30);
constexpr ll INFLL = 1LL << 62;
constexpr long double EPS = 1e-12;
constexpr ll MOD = (ll)((1E+9) + 7);

class Modulo {
private:
  ll MAX = MOD;
  vll fac_v, finv_v, inv_v;

public:
  void Modulo_setting(ll MAX_) {
    MAX = MAX_;

    fac_v.resize(MAX + 1);
    finv_v.resize(MAX + 1);
    inv_v.resize(MAX + 1);

    fac_v[0] = fac_v[1] = 1;
    finv_v[0] = finv_v[1] = 1;
    inv_v[1] = 1;
    for (int i = 2; i <= MAX; ++i) {
      // i!
      fac_v[i] = (fac_v[i - 1] * i) % MOD;
      // i^{-1}
      inv_v[i] = (MOD - (inv_v[MOD % i] * (MOD / i)) % MOD) % MOD;
      //(i!)^{-1}
      finv_v[i] = (finv_v[i - 1] * inv_v[i]) % MOD;
    }
  };

  Modulo(ll MAX_) {
    Modulo_setting(MAX_);
    return;
  };

  Modulo(void) {
    Modulo_setting(MOD);
    return;
  };

  ~Modulo(void) { return; };

  ll inv(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return inv_v[n];
  };

  ll fac(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return fac_v[n];
  };

  ll inv_fac(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return finv_v[n];
  };

  ll ModCom(ll n, ll k) {
    if (n < k)
      return 0;
    if (n < 0 || k < 0 || n - k < 0)
      return 0;
    return (fac_v[n] * ((finv_v[k] * finv_v[n - k]) % MOD)) % MOD;
  };

  ll ExCom(ll n, ll k) {
    if (n < k)
      return 0;
    if (n < 0 || k < 0)
      return 0;
    ll ans = 1;
    k = min(k, n - k);
    for (ll i = 0; i < k; ++i) {
      ans *= (n - i);
      ans /= (i + 1);
    }
    return ans;
  };

  ll pow(ll x, ll n) {
    ll res = 1;
    while (n > 0) {
      if (n & 1)
        res = res * x % MOD;
      x = x * x % MOD;
      n >>= 1;
    }
    return res;
  };

  unordered_map<ll, ll> prime_decomposer(ll n) {
    unordered_map<ll, ll> mp;
    for (ll i = 2; i * i <= n; ++i) {
      while (n % i == 0) {
        mp[i]++;
        n /= i;
      }
    }
    if (n != 1)
      mp[n]++;
    return mp;
  };

  ll LCM(vector<ll> vec) {
    unordered_map<ll, ll> mp;
    for (ll i = 0; i < vec.size(); ++i) {
      unordered_map<ll, ll> mp_tmp = prime_decomposer(vec[i]);
      for (auto e : mp_tmp) {
        mp[e.first] = max(mp[e.first], e.second);
      }
    }
    ll lcm = 1;
    for (auto e : mp) {
      lcm = (lcm * pow(e.first, e.second)) % MOD;
    }
    return lcm;
  };

  ll GCM(vector<ll> vec) {
    unordered_map<ll, ll> mp;
    for (ll i = 0; i < vec.size(); ++i) {
      unordered_map<ll, ll> mp_tmp = prime_decomposer(vec[i]);
      for (auto e : mp_tmp) {
        if (mp.find(e.first) == mp.end())
          mp[e.first] = e.second;
        else
          mp[e.first] = min(mp[e.first], e.second);
      }
    }
    ll gcm = 1;
    for (auto e : mp) {
      gcm = (gcm * pow(e.first, e.second)) % MOD;
    }
    return gcm;
  };
};

template <typename Data,
          typename T> // Data: the datatype of a data on each bond, T: the
                      // datatype of the data that each node has
                      class ReRooting {
private:
  struct Edge { // g[u][i].to = j <=> g[j][g[u][i].rev] = u
    ll to, rev;
    Data data;
  };

  using F1 =
      function<T(T, T)>; // the function to compare two data that each node has
  using F2 = function<T(T, Data)>; // the function to obtain the data on child
                                   // node through the bond

  ll N;
  vector<vector<Edge>> g;
  vector<vector<T>> ldp, rdp;
  vll lptr, rptr;
  const F1 f1;
  const F2 f2;

  T dfs(ll idx, ll par) {
    while (lptr[idx] != par && lptr[idx] < g[idx].size()) {
      auto &e = g[idx][lptr[idx]];
      ldp[idx][lptr[idx] + 1] =
          f1(ldp[idx][lptr[idx]], f2(dfs(e.to, e.rev), e.data));
      ++lptr[idx];
    }
    while (rptr[idx] != par && rptr[idx] >= 0) {
      auto &e = g[idx][rptr[idx]];
      rdp[idx][rptr[idx]] =
          f1(rdp[idx][rptr[idx] + 1], f2(dfs(e.to, e.rev), e.data));
      --rptr[idx];
    }
    if (par < 0)
      return rdp[idx][0];
    return f1(ldp[idx][par], rdp[idx][par + 1]);
  };

public:
  ReRooting(ll n, const F1 &f1, const F2 &f2)
      : N(n), g(n), ldp(n), rdp(n), lptr(n), rptr(n), f1(f1), f2(f2) {
    return;
  };

  ~ReRooting(void) { return; };

  void add_edge(ll u, ll v, const Data &d) {
    g[u].emplace_back((Edge){v, (ll)g[v].size(), d});
    g[v].emplace_back((Edge){u, (ll)g[u].size() - 1, d});
  };

  void add_edge_bi(ll u, ll v, const Data &d, const Data &e) {
    g[u].emplace_back((Edge){v, (ll)g[v].size(), d});
    g[v].emplace_back((Edge){u, (ll)g[u].size() - 1, e});
  };

  vector<T> solve(T init) {
    for (ll i = 0; i < N; ++i) {
      ldp[i].assign(g[i].size() + 1, init);
      rdp[i].assign(g[i].size() + 1, init);
      lptr[i] = 0;
      rptr[i] = (ll)g[i].size() - 1;
    }
    vector<T> ret(N);
    for (ll i = 0; i < N; ++i) {
      ret[i] = dfs(i, -1);
    }
    return ret;
  };

  vector<T> solve(vector<T> init) {
    for (ll i = 0; i < N; ++i) {
      ldp[i].assign(g[i].size() + 1, init[i]);
      rdp[i].assign(g[i].size() + 1, init[i]);
      lptr[i] = 0;
      rptr[i] = (ll)g[i].size() - 1;
    }
    vector<T> ret(N);
    for (ll i = 0; i < N; ++i) {
      ret[i] = dfs(i, -1);
    }
    return ret;
  };
};

Modulo que(1E6 + 7);

int main() {
  cin.tie(0); // cut the cin and cout (default, std::flush is performed after
              // std::cin)
  ios::sync_with_stdio(
      false); // cut the iostream and stdio (DON'T endl; BUT "\\n";)

  // pair<pll, pll> ans/prod of ans/# of node/prod of inv

  auto f1 = [](const pair<pll, pll> &a, const pair<pll, pll> &b) {
    pair<pll, pll> c = a;

    c.second.first += b.second.first;

    c.first.second *= b.first.second;
    c.first.second %= MOD;

    c.first.first = c.first.second * que.fac(c.second.first);
    c.first.first %= MOD;

    // debug_print(c.first.first, c.first.second, c.second.first,
    // c.second.second);
    return c;
  };

  auto f2 = [](const pair<pll, pll> &a, ll b) {
    pair<pll, pll> c = a;

    c.second.first += b;
    c.first.second = c.first.first * que.inv_fac(c.second.first);
    c.first.second %= MOD;

    return c;
  };

  ll N;
  IN(N);

  ReRooting<ll, pair<pll, pll>> g1(N, f1, f2);
  vvll G(N);

  REP(i, N - 1) {
    ll a, b;
    IN(a, b);
    a--;
    b--;
    G[a].pb(b);
    G[b].pb(a);
    g1.add_edge(a, b, 1);
  }

  pair<pll, pll> init;
  init.first.first = 1;
  init.first.second = 1;
  init.second.first = 0;
  init.second.second = 1;
  for (auto &p : g1.solve(init)) {
    OUT(p.first.first % MOD);
  }

  return 0;
}

Error: TLE

Fixed submission:

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

typedef long long ll;
typedef vector<ll> vll;
typedef vector<vector<ll>> vvll;
typedef vector<vector<vector<ll>>> vvvll;
typedef vector<bool> vb;
typedef vector<vector<bool>> vvb;
typedef vector<vector<vector<bool>>> vvvb;
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef vector<vpll> vvpll;
typedef vector<double> vd;
typedef vector<vd> vdd;
typedef pair<ll, vll> plvll;

#define FOR(i, x, y) for (ll i = (ll)x; i < (ll)y; ++i)
#define REP(i, y) FOR(i, 0, y)
#define RFOR(i, x, y) for (ll i = (ll)x; i >= (ll)y; --i)
#define RREP(i, x) RFOR(i, x, 0)
#define ALL(a) a.begin(), a.end()
#define pb push_back
#define debug_print(x...)                                                      \\
  cerr << "line " << __LINE__ << " : ";                                        \\
  debug_print_in(x);

template <typename First> void debug_print_in(First first) {
  cerr << first << endl;
  return;
}

template <typename First, typename... Rest>
void debug_print_in(First first, Rest... rest) {
  cerr << first << " ";
  debug_print_in(rest...);
  return;
}

void IN(void) { return; }

template <typename First, typename... Rest>
void IN(First &first, Rest &...rest) {
  cin >> first;
  IN(rest...);
  return;
}

template <typename First> void OUT(First first) {
  cout << first << endl;
  return;
}

template <typename First, typename... Rest>
void OUT(First first, Rest... rest) {
  cout << first << " ";
  OUT(rest...);
  return;
}

template <class t, class u> void chmax(t &a, u b) {
  if (a < b)
    a = b;
};
template <class t, class u> void chmin(t &a, u b) {
  if (a > b)
    a = b;
};
int popcount(int t) { return __builtin_popcount(t); }  // GCC
int popcount(ll t) { return __builtin_popcountll(t); } // GCC

template <typename T> void vec_print(vector<T> VEC) {
  REP(i, VEC.size()) { cerr << VEC[i] << " "; }
  cerr << endl;
};

template <typename T>
void mat_print(vector<vector<T>> MAT){
    REP(i, MAT.size()){REP(j, MAT[i].size()){cerr << MAT[i][j] << " ";
}
cerr << endl;
}
}
;

constexpr int INF = (1 << 30);
constexpr ll INFLL = 1LL << 62;
constexpr long double EPS = 1e-12;
constexpr ll MOD = (ll)((1E+9) + 7);

class Modulo {
private:
  ll MAX = MOD;
  vll fac_v, finv_v, inv_v;

public:
  void Modulo_setting(ll MAX_) {
    MAX = MAX_;

    fac_v.resize(MAX + 1);
    finv_v.resize(MAX + 1);
    inv_v.resize(MAX + 1);

    fac_v[0] = fac_v[1] = 1;
    finv_v[0] = finv_v[1] = 1;
    inv_v[1] = 1;
    for (int i = 2; i <= MAX; ++i) {
      // i!
      fac_v[i] = (fac_v[i - 1] * i) % MOD;
      // i^{-1}
      inv_v[i] = (MOD - (inv_v[MOD % i] * (MOD / i)) % MOD) % MOD;
      //(i!)^{-1}
      finv_v[i] = (finv_v[i - 1] * inv_v[i]) % MOD;
    }
  };

  Modulo(ll MAX_) {
    Modulo_setting(MAX_);
    return;
  };

  Modulo(void) {
    Modulo_setting(MOD);
    return;
  };

  ~Modulo(void) { return; };

  ll inv(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return inv_v[n];
  };

  ll fac(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return fac_v[n];
  };

  ll inv_fac(ll n) {
    if (n < 0 || n > MAX) {
      fprintf(stderr, "n range error\\n");
      exit(0);
    }
    return finv_v[n];
  };

  ll ModCom(ll n, ll k) {
    if (n < k)
      return 0;
    if (n < 0 || k < 0 || n - k < 0)
      return 0;
    return (fac_v[n] * ((finv_v[k] * finv_v[n - k]) % MOD)) % MOD;
  };

  ll ExCom(ll n, ll k) {
    if (n < k)
      return 0;
    if (n < 0 || k < 0)
      return 0;
    ll ans = 1;
    k = min(k, n - k);
    for (ll i = 0; i < k; ++i) {
      ans *= (n - i);
      ans /= (i + 1);
    }
    return ans;
  };

  ll pow(ll x, ll n) {
    ll res = 1;
    while (n > 0) {
      if (n & 1)
        res = res * x % MOD;
      x = x * x % MOD;
      n >>= 1;
    }
    return res;
  };

  unordered_map<ll, ll> prime_decomposer(ll n) {
    unordered_map<ll, ll> mp;
    for (ll i = 2; i * i <= n; ++i) {
      while (n % i == 0) {
        mp[i]++;
        n /= i;
      }
    }
    if (n != 1)
      mp[n]++;
    return mp;
  };

  ll LCM(vector<ll> vec) {
    unordered_map<ll, ll> mp;
    for (ll i = 0; i < vec.size(); ++i) {
      unordered_map<ll, ll> mp_tmp = prime_decomposer(vec[i]);
      for (auto e : mp_tmp) {
        mp[e.first] = max(mp[e.first], e.second);
      }
    }
    ll lcm = 1;
    for (auto e : mp) {
      lcm = (lcm * pow(e.first, e.second)) % MOD;
    }
    return lcm;
  };

  ll GCM(vector<ll> vec) {
    unordered_map<ll, ll> mp;
    for (ll i = 0; i < vec.size(); ++i) {
      unordered_map<ll, ll> mp_tmp = prime_decomposer(vec[i]);
      for (auto e : mp_tmp) {
        if (mp.find(e.first) == mp.end())
          mp[e.first] = e.second;
        else
          mp[e.first] = min(mp[e.first], e.second);
      }
    }
    ll gcm = 1;
    for (auto e : mp) {
      gcm = (gcm * pow(e.first, e.second)) % MOD;
    }
    return gcm;
  };
};

template <typename Data,
          typename T> // Data: the datatype of a data on each bond, T: the
                      // datatype of the data that each node has
                      class ReRooting {
private:
  struct Edge { // g[u][i].to = j <=> g[j][g[u][i].rev] = u
    ll to, rev;
    Data data;
  };

  using F1 =
      function<T(T, T)>; // the function to compare two data that each node has
  using F2 = function<T(T, Data)>; // the function to obtain the data on child
                                   // node through the bond

  ll N;
  vector<vector<Edge>> g;
  vector<vector<T>> ldp, rdp;
  vll lptr, rptr;
  const F1 f1;
  const F2 f2;

  T dfs(ll idx, ll par) {
    while (lptr[idx] != par && lptr[idx] < g[idx].size()) {
      auto &e = g[idx][lptr[idx]];
      ldp[idx][lptr[idx] + 1] =
          f1(ldp[idx][lptr[idx]], f2(dfs(e.to, e.rev), e.data));
      ++lptr[idx];
    }
    while (rptr[idx] != par && rptr[idx] >= 0) {
      auto &e = g[idx][rptr[idx]];
      rdp[idx][rptr[idx]] =
          f1(rdp[idx][rptr[idx] + 1], f2(dfs(e.to, e.rev), e.data));
      --rptr[idx];
    }
    if (par < 0)
      return rdp[idx][0];
    return f1(ldp[idx][par], rdp[idx][par + 1]);
  };

public:
  ReRooting(ll n, const F1 &f1, const F2 &f2)
      : N(n), g(n), ldp(n), rdp(n), lptr(n), rptr(n), f1(f1), f2(f2) {
    return;
  };

  ~ReRooting(void) { return; };

  void add_edge(ll u, ll v, const Data &d) {
    g[u].emplace_back((Edge){v, (ll)g[v].size(), d});
    g[v].emplace_back((Edge){u, (ll)g[u].size() - 1, d});
  };

  void add_edge_bi(ll u, ll v, const Data &d, const Data &e) {
    g[u].emplace_back((Edge){v, (ll)g[v].size(), d});
    g[v].emplace_back((Edge){u, (ll)g[u].size() - 1, e});
  };

  vector<T> solve(T init) {
    for (ll i = 0; i < N; ++i) {
      ldp[i].assign(g[i].size() + 1, init);
      rdp[i].assign(g[i].size() + 1, init);
      lptr[i] = 0;
      rptr[i] = (ll)g[i].size() - 1;
    }
    vector<T> ret(N);
    for (ll i = 0; i < N; ++i) {
      ret[i] = dfs(i, -1);
    }
    return ret;
  };

  vector<T> solve(vector<T> init) {
    for (ll i = 0; i < N; ++i) {
      ldp[i].assign(g[i].size() + 1, init[i]);
      rdp[i].assign(g[i].size() + 1, init[i]);
      lptr[i] = 0;
      rptr[i] = (ll)g[i].size() - 1;
    }
    vector<T> ret(N);
    for (ll i = 0; i < N; ++i) {
      ret[i] = dfs(i, -1);
    }
    return ret;
  };
};

Modulo que(1E6 + 7);

int main() {
  cin.tie(0); // cut the cin and cout (default, std::flush is performed after
              // std::cin)
  ios::sync_with_stdio(
      false); // cut the iostream and stdio (DON'T endl; BUT "\\n";)

  // pair<ll, plvll> ans/prod of ans/# of nodes

  auto f1 = [](const pair<ll, plvll> &a, const pair<ll, plvll> &b) {
    pair<ll, plvll> c;
    c.second.first = (a.second.first * b.second.first) % MOD;
    c.first = c.second.first;

    ll A_size = a.second.second.size();
    ll B_size = b.second.second.size();
    ll total = 0;
    c.second.second.resize(A_size + B_size);
    REP(i, A_size) {
      c.second.second[i] = a.second.second[i];
      total += a.second.second[i];
      c.first *= que.inv_fac(a.second.second[i]);
      c.first %= MOD;
    }
    REP(i, B_size) {
      c.second.second[i + A_size] = b.second.second[i];
      total += b.second.second[i];
      c.first *= que.inv_fac(b.second.second[i]);
      c.first %= MOD;
    }
    c.first *= que.fac(total);
    c.first %= MOD;
    return c;
  };

  auto f2 = [](const pair<ll, plvll> &a, ll b) {
    pair<ll, plvll> c;
    c.first = a.first;
    c.second.first = a.first;
    ll total = 0;
    for (auto i : a.second.second)
      total += i;
    vll tmp(1, total + b);
    c.second.second = tmp;
    return c;
  };

  ll N;
  IN(N);

  ReRooting<ll, pair<ll, plvll>> g1(N, f1, f2);
  vvll G(N);

  REP(i, N - 1) {
    ll a, b;
    IN(a, b);
    a--;
    b--;
    G[a].pb(b);
    G[b].pb(a);
    g1.add_edge(a, b, 1);
  }

  pair<ll, plvll> init;
  init.first = 1;
  init.second.first = 1;
  vll init_vll;
  init.second.second = init_vll;
  for (auto &p : g1.solve(init)) {
    OUT(p.first % MOD);
  }

  return 0;
}