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

#ifdef MAIN == == == == == == == == == == == == == == == == == == == == == ==  \\
    == == == == == == == == == == == == == == == == == == == == == == == == == \\
    == == == == == == == == == == == == == == == == ==

+ library {
name:
  ushi_rerooting
}

+ library {
name:
  mint, url : "https://tqkoh.github.io/library/library/lib/util/mint.hpp",
              version : "2020-08-08"
}

+ library {
name:
  ncr, url : "https://tqkoh.github.io/library/library/lib/util/ncr.hpp",
             version : "2020-08-08"
}

int solve(lint n, std::vector<lint> a, std::vector<lint> b) {
  auto merge = [&](P a, P b) { // choutensuu, toorisuu
    return P(a.f + b.f, (c.ncr(a.f + b.f, a.f) * a.s * b.s).a);
  };
  auto addroot = [&](P a, lint b) { return P(a.f + 1, a.s); };
  ReRooting<P, lint> g(n, merge, addroot, P(0, 1));
  rep(n - 1) g.add_edge(--a[i], --b[i], 1);
  auto dp = g.build();
  each(dp) out(e.s);
  return 0;
}

#endif MAIN == == == == == == == == == == == == == == == == == == == == == ==  \\
    == == == == == == == == == == == == == == == == == == == == == == == == == \\
    == == == == == == == == == == == == == == == == ==

// sub-BOF

// desktop
// author: Tqk

// #define _AOJ_
#define _C_INPUT_

// #pragma region template

// #pragma region basic
#define IN_FILE "./in.txt"
// #pragma GCC optimize ("O3")
#pragma warning(disable : 4455 4244 4067 4068 4996)
#pragma GCC target("avx")
#pragma GCC diagnostic ignored "-Wliteral-suffix"
#define NOMINMAX
#ifdef _ENV_TQK_
#include <Windows.h>
#endif
#include "bits/stdc++.h"
using namespace std;
typedef int64_t lint;
typedef long double ld;
typedef string cs;
#define linf 1152921504606846976
#ifndef _AOJ_
inline lint operator"" z(const unsigned long long n) { return lint(n); }
#endif
template <class T> inline T operator-(const vector<T> &a, const size_t &i) {
  return a.at(i);
}
// #pragma endregion

// #pragma region rep
#define _vcppunko3(tuple) _getname3 tuple
#define _vcppunko4(tuple) _getname4 tuple
#define _getname4(_1, _2, _3, _4, name, ...) name
#define _getname3(_1, _2, _3, name, ...) name
#define _trep2(tuple) _rep2 tuple
#define _trep3(tuple) _rep3 tuple
#define _trep4(tuple) _rep4 tuple
#define _rep1(n) for (lint i = 0; i < n; ++i)
#define _rep2(i, n) for (lint i = 0; i < n; ++i)
#define _rep3(i, a, b) for (lint i = a; i < b; ++i)
#define _rep4(i, a, b, c) for (lint i = a; i < b; i += c)
#define _tper2(tuple) _per2 tuple
#define _tper3(tuple) _per3 tuple
#define _tper4(tuple) _per4 tuple
#define _per1(n) for (lint i = n - 1; i >= 0; --i)
#define _per2(i, n) for (lint i = n - 1; i >= 0; --i)
#define _per3(i, a, b) for (lint i = b - 1; i >= a; --i)
#define _per4(i, a, b, c) for (lint i = a + (b - a - 1) / c * c; i >= a; i -= c)
#define rep(...)                                                               \\
  _vcppunko4((__VA_ARGS__, _trep4, _trep3, _trep2, _rep1))((__VA_ARGS__))
#define per(...)                                                               \\
  _vcppunko4((__VA_ARGS__, _tper4, _tper3, _tper2, _per1))((__VA_ARGS__))
#define _tall2(tuple) _all2 tuple
#define _tall3(tuple) _all3 tuple
#define _all1(v) v.begin(), v.end()
#define _all2(v, a) v.begin(), v.begin() + a
#define _all3(v, a, b) v.begin() + a, v.begin + b
#define all(...) _vcppunko3((__VA_ARGS__, _tall3, _tall2, _all1))((__VA_ARGS__))
#define each(c) for (auto &e : c)
// #pragma endregion

// #pragma region io
template <class T> istream &operator>>(istream &is, vector<T> &vec);
template <class T, size_t size>
istream &operator>>(istream &is, array<T, size> &vec);
template <class T, class L> istream &operator>>(istream &is, pair<T, L> &p);
template <class T> ostream &operator<<(ostream &os, vector<T> &vec);
template <class T, class L> ostream &operator<<(ostream &os, pair<T, L> &p);
template <class T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &x : vec)
    is >> x;
  return is;
}
template <class T, class L> istream &operator>>(istream &is, pair<T, L> &p) {
  is >> p.first;
  is >> p.second;
  return is;
}
template <class T, class L> ostream &operator<<(ostream &os, pair<T, L> &p) {
  os << p.first << " " << p.second;
  return os;
}
template <class T> ostream &operator<<(ostream &os, vector<T> &vec) {
  os << vec[0];
  rep(i, 1, vec.size()) os << ' ' << vec[i];
  return os;
}
template <class T> ostream &operator<<(ostream &os, deque<T> &deq) {
  os << deq[0];
  rep(i, 1, deq.size()) os << ' ' << deq[i];
  return os;
}
template <class T> ostream &operator<<(ostream &os, set<T> &s) {
  each(s) os << e << " ";
  return os;
}

#ifdef _ENV_TQK_
/*
ifstream infile(IN_FILE);
template<class T>
ifstream& operator>>(ifstream& is,vector<T>& vec);
template<class T,size_t size>
ifstream& operator>>(ifstream& is,array<T,size>& vec);
template<class T,class L>
ifstream& operator>>(ifstream& is,pair<T,L>& p);
template<class T>
ifstream& operator>>(ifstream& is,vector<T>& vec){ for(T& x: vec) is>>x;return
is; } template<class T,class L> ifstream& operator>>(ifstream& is,pair<T,L>& p){
is>>p.first;is>>p.second;return is; } inline void fin(){} template <class
Head,class... Tail> inline void fin(Head&& head,Tail&&... tail){
infile>>head;fin(move(tail)...); }
*/
// #include<Windows.h>
HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
inline void in() { SetConsoleTextAttribute(hConsole, 10); }
template <class Head, class... Tail>
inline void in(Head &&head, Tail &&...tail) {
  SetConsoleTextAttribute(hConsole, 15);
  cin >> head;
  in(move(tail)...);
}
#else
inline void in() {}
template <class Head, class... Tail>
inline void in(Head &&head, Tail &&...tail) {
  cin >> head;
  in(move(tail)...);
}
#endif

inline bool out() { return (cout << '\\n', 0); }
template <class T> inline bool out(T t) { return (cout << t << '\\n', 0); }
template <class Head, class... Tail> inline bool out(Head head, Tail... tail) {
  cout << head << ' ';
  out(move(tail)...);
  return 0;
}
template <class T> inline void alloc(T &c, lint s) {
  rep(c.size()) c[i].resize(s);
}
#define alc alloc
#ifdef _ENV_TQK_
inline bool deb() {
  SetConsoleTextAttribute(hConsole, 10);
  return (cout << '\\n', 0);
}
template <class T> inline bool deb(T t) {
  return (SetConsoleTextAttribute(hConsole, 12), cout << t << '\\n',
          SetConsoleTextAttribute(hConsole, 10), 0);
}
template <class Head, class... Tail> inline bool deb(Head head, Tail... tail) {
  SetConsoleTextAttribute(hConsole, 12);
  cout << head << ' ';
  deb(move(tail)...);
  return 0;
}
#define dsp(ex) sp(ex)
#define dno(ex) no(ex)
#define look(v)                                                                \\
  SetConsoleTextAttribute(hConsole, 12), cout << #v << ": ", deb(v);
#else
#define deb(...) 0
#define dsp(ex) 0
#define dno(ex) 0
#define look(v) 0
#endif

// #pragma endregion

// #pragma region TA
#define lin(...)                                                               \\
  lint __VA_ARGS__;                                                            \\
  in(__VA_ARGS__)
#define stin(...)                                                              \\
  string __VA_ARGS__;                                                          \\
  in(__VA_ARGS__)
#define vin(type, name, size)                                                  \\
  vector<type> name(size);                                                     \\
  in(name)
#define pb push_back
#define fi e.first
#define se e.second
#define YES(c) cout << ((c) ? "YES\\n" : "NO\\n"), 0
#define Yes(c) cout << ((c) ? "Yes\\n" : "No\\n"), 0
#define POSSIBLE(c) cout << ((c) ? "POSSIBLE\\n" : "IMPOSSIBLE\\n"), 0
#define Possible(c) cout << ((c) ? "Possible\\n" : "Impossible\\n"), 0
#define o(p) cout << p << endl, 0
#define sp(p) cout << p << " ", 0
#define no(p) cout << p, 0
#define st(v) sort(all(v))
#define psort(l, r)                                                            \\
  if (l > r)                                                                   \\
    swap(l, r);
#define fn(ty1, ty2, ex) [&](ty1 l, ty2 r) { return (ex); }
#define lfn(ex) [&](lint l, lint r) { return (ex); }

// #pragma endregion

// #pragma region func
inline constexpr lint gcd(lint a, lint b) {
  while (b) {
    lint c = b;
    b = a % b;
    a = c;
  }
  return a;
}
inline constexpr lint lcm(lint a, lint b) { return a / gcd(a, b) * b; }
template <typename T> inline constexpr bool chmin(T &mn, const T &cnt) {
  if (mn > cnt) {
    mn = cnt;
    return 1;
  } else
    return 0;
}
template <typename T> inline constexpr bool chmax(T &mx, const T &cnt) {
  if (mx < cnt) {
    mx = cnt;
    return 1;
  } else
    return 0;
}
// #pragma endregion

// #pragma region P

class P {
public:
  lint f, s;
  P(lint a, lint b) : f(a), s(b){};
  P() : f(0), s(0){};
};
istream &operator>>(istream &os, P &p) {
  os >> p.f >> p.s;
  return os;
}
inline constexpr bool operator<(const P &l, const P &r) {
  return (l.f - r.f ? l.f < r.f : l.s < r.s);
}
inline constexpr bool operator>(const P &l, const P &r) {
  return (l.f - r.f ? l.f > r.f : l.s > r.s);
}
inline constexpr bool operator==(const P &l, const P &r) {
  return (l.f == r.f && l.s == r.s);
}
inline constexpr bool operator!=(const P &l, const P &r) {
  return (l.f != r.f || l.s != r.s);
}
inline P operator+(const P &l, const P &r) { return P(l.f + r.f, l.s + r.s); }
inline P operator-(const P &l, const P &r) { return P(l.f - r.f, l.s - r.s); }
inline P operator-(const P &r) { return P(-r.f, -r.s); }
inline P operator*(const lint &l, const P &r) { return P(l * r.f, l * r.s); }
inline P operator/(const P &l, const lint &r) { return P(l.f / r, l.s / r); }
inline ld abs(P p) { return sqrtl(p.f * p.f + p.s * p.s); }
inline ld mabs(P p) { return abs(p.f) + abs(p.s); }
inline ld mht(P x, P y) { return mabs(x - y); }

// #pragma endregion

// #pragma region start
struct Start {
  Start() {
#ifndef _C_INPUT_
    cin.tie(0), cout.tie(0);
    ios::sync_with_stdio(0);
#endif
    cout << fixed << setprecision(16);
  }
} __start;
// #pragma endregion

// #pragma endregion

// #pragma region const
#define mod 1000000007
const ld pi = acos(-1);
const ld tau = (1. + sqrtl(5)) / 2.;
P d4[4] = {P(1, 0), P(0, 1), P(-1, 0), P(0, -1)};
P d8[8] = {P(1, 0),  P(1, 1),   P(0, 1),  P(-1, 1),
           P(-1, 0), P(-1, -1), P(0, -1), P(1, -1)};
const cs AUTO_YES = "Yes";
const cs AUTO_NO = "No";
int cho(bool c, cs yes = AUTO_YES, cs no = AUTO_NO) {
  return out((c ? yes : no));
}
// #pragma endregion

// #pragma region solve

// #pragma region lib_ushi_rerooting

template <typename sum_t, typename key_t> struct ReRooting {
  struct Edge {
    int to;
    key_t data;
    sum_t dp, ndp;
  };

  using F = function<sum_t(sum_t, sum_t)>;
  using G = function<sum_t(sum_t, key_t)>;

  vector<vector<Edge>> g;
  vector<sum_t> subdp, dp;
  const sum_t ident;
  const F f;
  const G gg;

  ReRooting(int V, const F f, const G g, const sum_t &ident)
      : g(V), f(f), gg(g), ident(ident), subdp(V, ident), dp(V, ident) {}

  void add_edge(int u, int v, const key_t &d) {
    g[u].push_back({v, d, ident, ident});
    g[v].push_back({u, d, ident, ident});
  }

  void add_edge_bi(int u, int v, const key_t &d, const key_t &e) {
    g[u].emplace_back((Edge){v, d, ident, ident});
    g[v].emplace_back((Edge){u, e, ident, ident});
  }

  void dfs_sub(int idx, int par) {
    for (auto &e : g[idx]) {
      if (e.to == par)
        continue;
      dfs_sub(e.to, idx);
      subdp[idx] = f(subdp[idx], gg(subdp[e.to], e.data));
    }
  }

  void dfs_all(int idx, int par, const sum_t &top) {
    sum_t buff{ident};
    for (int i = 0; i < (int)g[idx].size(); i++) {
      auto &e = g[idx][i];
      e.ndp = buff;
      e.dp = gg(par == e.to ? top : subdp[e.to], e.data);
      buff = f(buff, e.dp);
    }
    dp[idx] = buff;
    buff = ident;
    for (int i = (int)g[idx].size() - 1; i >= 0; i--) {
      auto &e = g[idx][i];
      if (e.to != par)
        dfs_all(e.to, idx, f(e.ndp, buff));
      e.ndp = f(e.ndp, buff);
      buff = f(buff, e.dp);
    }
  }

  vector<sum_t> build() {
    dfs_sub(0, -1);
    dfs_all(0, -1, ident);
    return dp;
  }
};

// #pragma endregion

// #pragma region lib_mint

#define md_tmp template <uint_fast64_t md = mod>
md_tmp class modint {
  using u64 = uint_fast64_t;

public:
  u64 a;

  constexpr modint(const lint x = 0) noexcept : a((x + md) % md) {}
  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 lint 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 >= md)
      a -= md;
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a)
      a += md;
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % md;
    return *this;
  }
  constexpr modint &operator^=(const lint rhs) noexcept {
    if (!rhs)
      return *this = 1;
    u64 exp = rhs - 1;
    modint base = a;
    while (exp) {
      if (exp & 1)
        *this *= base;
      base *= base;
      exp >>= 1;
    }
    return *this;
  }
  constexpr modint &operator/=(const modint rhs) noexcept {
    a = (*this * (rhs ^ (md - 2))).a;
    return *this;
  }
};
md_tmp istream &operator>>(istream &os, modint<md> &m) {
  os >> m.a, m.a %= md;
  return os;
}
md_tmp ostream &operator<<(ostream &os, const modint<md> &m) {
  return os << m.a;
}
using mint = modint<>;
#ifndef _AOJ_
mint operator"" m(unsigned long long n) { return mint(n); }
#endif

/*
 * @title Mint
 * @see https://noshi91.hatenablog.com/entry/2019/03/31/174006
 */

// #pragma endregion

// #pragma region lib_ncr

#define MAXN 200010
md_tmp class Ncr {
  lint m;

public:
  mint fa[MAXN], rfa[MAXN];
  Ncr(lint n) : m(n) {
    fa[0] = 1;
    for (int i = 1; i < m; ++i)
      fa[i] = (fa[i - 1] * i);
    rfa[m - 1] = mint(1) / fa[m - 1];
    for (int i = m - 2; i >= 0; --i)
      rfa[i] = rfa[i + 1] * (i + 1);
  }
  inline mint ncr(lint n, lint r) {
    if (n < r || n < 0 || r < 0)
      return 0m;
    return (fa[n] * rfa[r] * rfa[n - r]);
  }
};
Ncr<> c(MAXN);
#undef MAXN

/*
 * @title Ncr
 */

// #pragma endregion

int solve(lint n, std::vector<lint> a, std::vector<lint> b) {
  auto merge = [&](P a, P b) { // choutensuu, toorisuu
    return P(a.f + b.f, (c.ncr(a.f + b.f, a.f) * a.s * b.s).a);
  };
  auto addroot = [&](P a, lint b) { return P(a.f + 1, a.s); };
  ReRooting<P, lint> g(n, merge, addroot, P(0, 1));
  rep(n - 1) g.add_edge(--a[i], --b[i], 1);
  auto dp = g.build();
  each(dp) out(e.s);
  return 0;
}

// #pragma endregion

// #pragma region main

int main() {
  lint N;
  scanf("%lld", &N);
  std::vector<lint> a(N - 1);
  std::vector<lint> 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;
}

// #pragma endregion

// sub-EOF


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <atomic>
#include <bitset>
#include <cassert>
#include <ccomplex>
#include <cctype>
#include <cerrno>
#include <cfenv>
#include <cfloat>
#include <chrono>
#include <cinttypes>
#include <ciso646>
#include <climits>
#include <clocale>
#include <cmath>
#include <complex>
#include <condition_variable>
#include <csetjmp>
#include <csignal>
#include <cstdarg>
#include <cstdbool>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctgmath>
#include <ctime>
#include <cwchar>
#include <cwctype>
#include <deque>
#include <exception>
#include <forward_list>
#include <fstream>
#include <functional>
#include <future>
#include <initializer_list>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <limits>
#include <list>
#include <locale>
#include <map>
#include <memory>
#include <mutex>
#include <new>
#include <numeric>
#include <ostream>
#include <queue>
#include <random>
#include <ratio>
#include <regex>
#include <set>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <streambuf>
#include <string>
#include <system_error>
#include <thread>
#include <tuple>
#include <type_traits>
#include <typeindex>
#include <typeinfo>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <valarray>
#include <vector>

using namespace std;

const long long mod = 1000000007;

struct Combination {
  int mod;
  vector<int64_t> mfact, rfact;

  Combination(int sz, int mod) : mfact(sz + 1), rfact(sz + 1), mod(mod) {
    mfact[0] = 1;
    for (int i = 1; i < mfact.size(); i++) {
      mfact[i] = mfact[i - 1] * i % mod;
    }
    rfact[sz] = inv(mfact[sz]);
    for (int i = sz - 1; i >= 0; i--) {
      rfact[i] = rfact[i + 1] * (i + 1) % mod;
    }
  }

  int64_t fact(int k) const { return (mfact[k]); }

  int64_t pow(int64_t x, int64_t n) const {
    int64_t ret = 1;
    while (n > 0) {
      if (n & 1)
        (ret *= x) %= mod;
      (x *= x) %= mod;
      n >>= 1;
    }
    return (ret);
  }

  int64_t inv(int64_t x) const { return (pow(x, mod - 2)); }

  int64_t P(int n, int r) const {
    if (r < 0 || n < r)
      return (0);
    return (mfact[n] * rfact[n - r] % mod);
  }

  int64_t C(int p, int q) const {
    if (q < 0 || p < q)
      return (0);
    return (mfact[p] * rfact[q] % mod * rfact[p - q] % mod);
  }

  int64_t H(int n, int r) const {
    if (n < 0 || r < 0)
      return (0);
    return (r == 0 ? 1 : C(n + r - 1, r));
  }
};

struct phash {
  inline long long operator()(const pair<long long, long long> &p) const {
    const auto h1 = hash<long long>()(p.first);
    const auto h2 = hash<long long>()(p.second);
    return h1 ^ (h2 << 1);
  }
};

Combination comb(202020, mod);

unordered_map<pair<long long, long long>, pair<long long, long long>, phash> mp;

vector<vector<int>> G(202020);

pair<long long, long long> f(int p, int c) {
  if (mp[pair<int, int>(p, c)].first)
    return mp[pair<long long, long long>(p, c)];
  if (G[c].size() == 1 && G[c][0] == p) {
    mp[make_pair(p, c)].first = 1;
    mp[make_pair(p, c)].second = 1;
    return mp[make_pair(p, c)];
  }

  long long ret = 1;
  long long sum = 0;
  for (int i = 0; i < G[c].size(); i++) {
    int g = G[c][i];
    if (g == p)
      continue;
    pair<long long, long long> d = f(c, g);
    ret *= d.first;
    sum += d.second;
    ret %= mod;

    ret *= comb.inv(comb.fact(d.second));
    ret %= mod;
  }
  // cout << sum << endl;
  ret = ret * comb.fact(sum);
  ret %= mod;
  return mp[pair<int, int>(p, c)] = pair<long long, long long>(ret, sum + 1);
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

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

  //   f(N + 1, 0);

  for (int i = 0; i < N; i++) {
    cout << f(N + 1, i).first << endl;
  }
  /*        cout << mp[make_pair(0, 1)].first << endl;
          cout << mp[make_pair(0, 1)].second << endl;*/
}