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

#include <algorithm>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <deque>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define FOR(i, k, n) for (int(i) = (k); (i) < (n); ++(i))
#define rep(i, n) FOR(i, 0, n)
#define all(v) begin(v), end(v)
#define debug(x)     // std::cerr<<#x<<": "<<x<<"\\n"
#define debug2(x, y) // std::cerr<<#x<<": "<<x<<", "<<#y<<": "<<y<<"\\n"
#define debug3(                                                                \\
    x, y,                                                                      \\
    z) // std::cerr<<#x<<": "<<x<<", "<<#y<<": "<<y<<", "<<#z<<": "<<z<<"\\n"

using ll = long long;
using vi = std::vector<int>;
using vvi = std::vector<vi>;
using vll = std::vector<ll>;
using vvll = std::vector<vll>;
template <typename T> using vvec = std::vector<std::vector<T>>;
template <typename T> auto make_v(size_t sz) { return std::vector<T>(sz); }
template <typename T, typename... Ts> auto make_v(size_t sz, Ts... ts) {
  return std::vector<decltype(make_v<T>(ts...))>(sz, make_v<T>(ts...));
}
template <typename T> void fill_v(T &var, const T &x) { var = x; }
template <typename V, typename T> void fill_v(V &v, const T &x) {
  for (auto &&w : v) {
    fill_v(w, x);
  }
}

template <typename T>
std::ostream &operator<<(std::ostream &s, const std::vector<T> &v) {
  int sz = v.size();
  s << "\\n";
  rep(i, sz) {
    s << v[i];
    if (i < sz - 1) {
      s << "\\t";
    }
  }
  s << "\\n";
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s,
                         const std::vector<std::vector<T>> &v) {
  for (auto &&w : v) {
    s << w;
  }
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::deque<T> &v) {
  int sz = v.size();
  s << "\\n";
  rep(i, sz) {
    s << v[i];
    if (i < sz - 1) {
      s << "\\t";
    }
  }
  s << "\\n";
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::deque<std::deque<T>> &v) {
  for (auto &&w : v) {
    s << w;
  }
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::set<T> &v) {
  s << "\\n";
  for (auto &&elm : v) {
    s << elm << "\\t";
  }
  s << "\\n";
  return s;
}

inline void scan(int &a) { scanf("%d", &a); }
inline void scan(ll &a) { scanf("%lld", &a); }
inline void scan(char &a) { scanf(" %c", &a); }
inline void scan(double &a) { scanf("%lf", &a); }
template <typename T> inline void scan(std::vector<T> &v) {
  for (auto &&sv : v) {
    scan(sv);
  }
}
template <typename First, typename... Args>
inline void scan(First &f, Args &...args) {
  scan(f);
  scan(args...);
}
inline void scan(std::string &s) {
  char BUF[3000000];
  scanf(" %s", BUF);
  s = std::string(BUF);
}
inline void print(int a) { printf("%d\\n", a); }
inline void print(ll a) { printf("%lld\\n", a); }
inline void print(double a) { printf("%.12f\\n", a); }
inline void print(std::string s) { std::cout << s << "\\n"; }

using namespace std;

int mod = 1000000007;
ll po(ll k, ll x) {
  if (x == 0) {
    return 1;
  }
  if (x == 1) {
    return k % mod;
  }
  ll y = po(k, x / 2);
  y = y * y % mod;
  if (x % 2 == 1) {
    y = y * k % mod;
  }
  return y;
}
vll fac, inv;
ll combination(ll n, ll k) { return fac[n] * inv[k] % mod * inv[n - k] % mod; }

int n;
void print_(vll &ans) {
  rep(i, n) { print(ans[i]); }
  exit(0);
}

vvi g;
vll ans;
vector<unordered_map<int, ll>> pans;  // answer for partial tree
vector<unordered_map<int, ll>> pvnum; // number of vertices in partial tree

vll calc_aux(int proot, int rootpar) {
  bool pans_calced = (pans[proot][rootpar] != -1);
  bool pvnum_calced = (pvnum[proot][rootpar] != -1);
  assert(!(pans_calced ^ pvnum_calced));
  if (pans_calced and pvnum_calced) {
    return {pans[proot][rootpar], pvnum[proot][rootpar]};
  }
  ll ret = 1;
  vll pvnums;
  int sz = g[proot].size() - 1;
  pvnums.reserve(sz);
  ll pvnum_sum = 0;
  for (int u : g[proot]) {
    if (u == rootpar) {
      continue;
    }
    auto pret = calc_aux(u, proot);
    ret = ret * pret[0] % mod;
    pvnums.push_back(pret[1]);
    pvnum_sum += pret[1];
  }
  assert((int)pvnums.size() == sz);
  ll csum = pvnum_sum;
  rep(i, sz) {
    ret = ret * combination(csum, pvnums[i]) % mod;
    csum -= pvnums[i];
  }
  pans[proot][rootpar] = ret;
  pvnum[proot][rootpar] = pvnum_sum + 1;
  return {ret, pvnum_sum + 1};
}

void calc_aux2(int proot, int rootpar) {
  debug3(ans[rootpar], pans[proot][rootpar], pvnum[proot][rootpar]);
  assert((ans[rootpar] != -1) and (pans[proot][rootpar] != -1) and
         (pvnum[proot][rootpar] != -1));
  ll inv_pans = po(pans[proot][rootpar], mod - 2);
  ll inv_comb = po(combination(n - 1, pvnum[proot][rootpar]), mod - 2);
  ll pans_rootpar = ans[rootpar] * inv_pans % mod * inv_comb % mod;
  debug3(inv_pans, inv_comb, pans_rootpar);
  ans[proot] = pans[proot][rootpar] * pans_rootpar % mod *
               combination(n - 1, pvnum[proot][rootpar] - 1) % mod;
  for (int u : g[proot]) {
    if (u == rootpar) {
      continue;
    }
    calc_aux2(u, proot);
  }
}

void calc(int root) {
  if (ans[root] != -1) {
    return;
  }
  ll ret = 1;
  vll pvnums;
  int sz = g[root].size();
  pvnums.reserve(sz);
  ll pvnum_sum = 0;
  for (int u : g[root]) {
    auto pret = calc_aux(u, root);
    ret = ret * pret[0] % mod;
    pvnums.push_back(pret[1]);
    pvnum_sum += pret[1];
  }
  assert((int)pvnums.size() == sz);
  assert(pvnum_sum == n - 1);
  ll csum = pvnum_sum;
  rep(i, sz) {
    ret = ret * combination(csum, pvnums[i]) % mod;
    csum -= pvnums[i];
  }
  ans[root] = ret;
  for (int u : g[root]) {
    calc_aux2(u, root);
  }
}

int main() {
  scan(n);
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    scan(a, b);
    --a;
    --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }

  fac.assign(n + 1, 0);
  inv.assign(n + 1, 0);
  fac[0] = 1;
  inv[0] = 1;
  rep(i, n) {
    fac[i + 1] = fac[i] * (i + 1) % mod;
    inv[i + 1] = po(fac[i + 1], mod - 2);
  }

  ans.assign(n, -1);
  pans.resize(n);
  pvnum.resize(n);
  rep(u, n) {
    for (int v : g[u]) {
      pans[u][v] = -1;
      pvnum[u][v] = -1;
    }
  }
  calc(0);
  print_(ans);

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <deque>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define FOR(i, k, n) for (int(i) = (k); (i) < (n); ++(i))
#define rep(i, n) FOR(i, 0, n)
#define all(v) begin(v), end(v)
#define debug(x) std::cerr << #x << ": " << x << "\\n"
#define debug2(x, y)                                                           \\
  std::cerr << #x << ": " << x << ", " << #y << ": " << y << "\\n"
#define debug3(x, y, z)                                                        \\
  std::cerr << #x << ": " << x << ", " << #y << ": " << y << ", " << #z        \\
            << ": " << z << "\\n"

using ll = long long;
using vi = std::vector<int>;
using vvi = std::vector<vi>;
using vll = std::vector<ll>;
using vvll = std::vector<vll>;
template <typename T> using vvec = std::vector<std::vector<T>>;
template <typename T> auto make_v(size_t sz) { return std::vector<T>(sz); }
template <typename T, typename... Ts> auto make_v(size_t sz, Ts... ts) {
  return std::vector<decltype(make_v<T>(ts...))>(sz, make_v<T>(ts...));
}
template <typename T> void fill_v(T &var, const T &x) { var = x; }
template <typename V, typename T> void fill_v(V &v, const T &x) {
  for (auto &&w : v) {
    fill_v(w, x);
  }
}

template <typename T>
std::ostream &operator<<(std::ostream &s, const std::vector<T> &v) {
  int sz = v.size();
  s << "\\n";
  rep(i, sz) {
    s << v[i];
    if (i < sz - 1) {
      s << "\\t";
    }
  }
  s << "\\n";
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s,
                         const std::vector<std::vector<T>> &v) {
  for (auto &&w : v) {
    s << w;
  }
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::deque<T> &v) {
  int sz = v.size();
  s << "\\n";
  rep(i, sz) {
    s << v[i];
    if (i < sz - 1) {
      s << "\\t";
    }
  }
  s << "\\n";
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::deque<std::deque<T>> &v) {
  for (auto &&w : v) {
    s << w;
  }
  return s;
}
template <typename T>
std::ostream &operator<<(std::ostream &s, const std::set<T> &v) {
  s << "\\n";
  for (auto &&elm : v) {
    s << elm << "\\t";
  }
  s << "\\n";
  return s;
}

inline void scan(int &a) { scanf("%d", &a); }
inline void scan(ll &a) { scanf("%lld", &a); }
inline void scan(char &a) { scanf(" %c", &a); }
inline void scan(double &a) { scanf("%lf", &a); }
template <typename T> inline void scan(std::vector<T> &v) {
  for (auto &&sv : v) {
    scan(sv);
  }
}
template <typename First, typename... Args>
inline void scan(First &f, Args &...args) {
  scan(f);
  scan(args...);
}
inline void scan(std::string &s) {
  char BUF[3000000];
  scanf(" %s", BUF);
  s = std::string(BUF);
}
inline void print(int a) { printf("%d\\n", a); }
inline void print(ll a) { printf("%lld\\n", a); }
inline void print(double a) { printf("%.12f\\n", a); }
inline void print(std::string s) { std::cout << s << "\\n"; }

using namespace std;

int mod = 1000000007;
ll po(ll k, ll x) {
  if (x == 0) {
    return 1;
  }
  if (x == 1) {
    return k % mod;
  }
  ll y = po(k, x / 2);
  y = y * y % mod;
  if (x % 2 == 1) {
    y = y * k % mod;
  }
  return y;
}
vll fac, inv;
ll combination(ll n, ll k) { return fac[n] * inv[k] % mod * inv[n - k] % mod; }

int n;
vvi g;
vll ans;
vector<unordered_map<int, ll>> pans;  // answer for partial tree
vector<unordered_map<int, ll>> pvnum; // number of vertices in partial tree

vll calc_aux(int proot, int rootpar) {
  bool pans_calced = (pans[proot][rootpar] != -1);
  bool pvnum_calced = (pvnum[proot][rootpar] != -1);
  assert(!(pans_calced ^ pvnum_calced));
  if (pans_calced and pvnum_calced) {
    return {pans[proot][rootpar], pvnum[proot][rootpar]};
  }
  ll ret = 1;
  vll pvnums;
  int sz = g[proot].size() - 1;
  pvnums.reserve(sz);
  ll pvnum_sum = 0;
  for (int u : g[proot]) {
    if (u == rootpar) {
      continue;
    }
    auto pret = calc_aux(u, proot);
    ret = ret * pret[0] % mod;
    pvnums.push_back(pret[1]);
    pvnum_sum += pret[1];
  }
  assert((int)pvnums.size() == sz);
  ll csum = pvnum_sum;
  rep(i, sz) {
    ret = ret * combination(csum, pvnums[i]) % mod;
    csum -= pvnums[i];
  }
  return {ret, pvnum_sum + 1};
}

ll calc(int root) {
  if (ans[root] != -1) {
    return ans[root];
  }
  ll ret = 1;
  vll pvnums;
  int sz = g[root].size();
  pvnums.reserve(sz);
  ll pvnum_sum = 0;
  for (int u : g[root]) {
    auto pret = calc_aux(u, root);
    ret = ret * pret[0] % mod;
    pvnums.push_back(pret[1]);
    pvnum_sum += pret[1];
  }
  assert((int)pvnums.size() == sz);
  assert(pvnum_sum == n - 1);
  ll csum = pvnum_sum;
  rep(i, sz) {
    ret = ret * combination(csum, pvnums[i]) % mod;
    csum -= pvnums[i];
  }
  return ret;
}

void print_(vll &ans) {
  rep(i, n) { print(ans[i]); }
  exit(0);
}

int main() {
  scan(n);
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    scan(a, b);
    --a;
    --b;
    g[a].push_back(b);
    g[b].push_back(a);
  }

  fac.assign(n + 1, 0);
  inv.assign(n + 1, 0);
  fac[0] = 1;
  inv[0] = 1;
  rep(i, n) {
    fac[i + 1] = fac[i] * (i + 1) % mod;
    inv[i + 1] = po(fac[i + 1], mod - 2);
  }

  ans.assign(n, -1);
  pans.resize(n);
  pvnum.resize(n);
  rep(u, n) {
    for (int v : g[u]) {
      pans[u][v] = -1;
      pvnum[u][v] = -1;
    }
  }
  rep(i, n) { ans[i] = calc(i); }
  print_(ans);

  return 0;
}