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

#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <numeric>
#include <string>
#include <utility>
#include <vector>
using namespace std;
using i8t = std::int64_t;
using u8t = std::uint64_t;
using vi8t = std::vector<i8t>;
using vvi8t = std::vector<std::vector<i8t>>;
using vu8t = std::vector<u8t>;
#include <limits>
static const i8t i8t_max = numeric_limits<i8t>::max();
static const i8t i8t_min = numeric_limits<i8t>::min();
#include <unordered_map>
#include <unordered_set>
template <typename Key, typename Value>
using map_t = std::unordered_map<Key, Value>;
template <typename Key> using set_t = std::unordered_set<Key>;
#define F3(i, i0, iN, step) for (i8t i = i0, i##_N = iN; i < i##_N; i += step)
#define F2(i, i0, iN) F3(i, i0, iN, 1)
#define F(i, N) F3(i, 0, N, 1)
#define FR(i, N) for (i8t i = N; i--;)
#define R(V)                                                                   \\
  i8t V;                                                                       \\
  std::cin >> V;
#define RV(V, N)                                                               \\
  std::vector<i8t> V;                                                          \\
  F(_i, N) { R(_x) V.push_back(_x); }
#define RL1(line) std::getline(std::cin, line);
#define RL(line)                                                               \\
  std::string line;                                                            \\
  RL1(line)
#define RM(V, L, C)                                                            \\
  std::vector<char> V;                                                         \\
  F(_i, L) { RL(_l) F(_j, C) V.push_back(_l[_j]); }
#define P(X) std::cout << (X) << std::endl;
template <typename T> T umax(T &ans, T const &value) {
  if (value > ans)
    ans = value;
  return ans;
}
template <typename T> T umin(T &ans, T const &value) {
  if (value < ans)
    ans = value;
  return ans;
}
template <typename T> struct identity {
  using type = T;
};

struct node_t {
  i8t weight = 0; // 子孫の数(自分も入れる)

  i8t parent = -1;
  std::vector<i8t> adj;

  i8t npat_child;
  i8t npat_parent;
};

void tree_make_directed_by_root(std::vector<node_t> &vtx, i8t root) {
  std::vector<bool> mark(vtx.size());
  std::vector<i8t> stk;
  std::vector<i8t> ind;
  i8t p, i;
  mark[p = root] = 1;
  i = 0;
  while (1) {
    node_t &node = vtx[p];
    if (i < node.adj.size()) {
      i8t &p2 = node.adj[i];
      if (mark[p2]) {
        p2 = -1;
        i++;
      } else {
        stk.push_back(p);
        ind.push_back(i + 1);
        mark[p = p2] = 1;
        i = 0;
      }
    } else {
      // remove parents
      node.adj.erase(std::remove(node.adj.begin(), node.adj.end(), -1),
                     node.adj.end());

      // update weight
      node.weight = 1;
      for (i8t p2 : node.adj)
        node.weight += vtx[p2].weight;

      if (stk.size()) {
        p = stk.back();
        i = ind.back();
        stk.pop_back();
        ind.pop_back();
      } else {
        break;
      }
    }
  }
}

template <typename FEnter, typename FLeave>
void tree_visit(std::vector<node_t> const &vtx, i8t root, FEnter _enter,
                FLeave _leave) {
  std::vector<i8t> stk, ind;
  i8t p, i;
  p = root;
  i = 0;
  _enter(p, -1);
  while (1) {
    auto &node = vtx[p];
    if (i < node.adj.size()) {
      stk.push_back(p);
      ind.push_back(i + 1);
      _enter(node.adj[i], p);
      p = node.adj[i];
      i = 0;
    } else {
      if (stk.size()) {
        _leave(p, stk.back());
        p = stk.back();
        i = ind.back();
        stk.pop_back();
        ind.pop_back();
      } else {
        _leave(p, -1);
        break;
      }
    }
  }
}

u8t u8t_gcd(u8t a, u8t b) {
  if (a < b)
    swap(a, b);
  if (b == 0)
    return 0;
  for (;;) {
    a %= b;
    if (a == 0)
      return b;
    b %= a;
    if (b == 0)
      return a;
  }
}

i8t i8t_gcd(i8t a, i8t b, i8t &s, i8t &t) {
  if (a == 0 || b == 0) {
    s = t = 0;
    return 0;
  } else if (a < 0) {
    i8t r = i8t_gcd(-a, b, s, t);
    s = -s;
    return r;
  } else if (b < 0) {
    i8t r = i8t_gcd(a, -b, s, t);
    t = -t;
    return r;
  }

  i8t q, as = 1, at = 0, bs = 0, bt = 1;
  for (;;) {
    q = a / b;
    a %= b;
    if (a == 0) {
      s = bs;
      t = bt;
      return b;
    }
    as -= q * bs;
    at -= q * bt;

    q = b / a;
    b %= a;
    if (b == 0) {
      s = as;
      t = at;
      return a;
    }
    bs -= q * as;
    bt -= q * at;
  }
}

/// @fn mod_inv(x,m)
///   互いに素な整数 x,m に対して x^{-1} mod m を求める。
/// @param[in] x
/// @param[in] m
i8t mod_inv(i8t x, i8t mod) {
  i8t xinv, b;
  if (1 != i8t_gcd(x, mod, xinv, b))
    std::exit(2); // 互いに素でない。
  xinv %= mod;
  if (xinv < 0)
    xinv += mod;
  return xinv;
}

i8t mod_div(i8t x, i8t y, i8t mod) { return x * mod_inv(y, mod) % mod; }

i8t mod_factorial(i8t n, i8t mod) {
  if (n <= 0)
    return n < 0 ? 0 : 1;
  i8t r = n;
  while (--n)
    r = r * n % mod;
  return r;
}

i8t mod_nCr(u8t n, u8t r, u8t mod) {
  i8t num = mod_factorial(n, mod);
  i8t d1 = mod_factorial(r, mod);
  i8t d2 = mod_factorial(n - r, mod);
  return mod_div(mod_div(num, d1, mod), d2, mod);
}

vi8t mod_factorial_table(i8t n, i8t mod) {
  if (n <= 0)
    return vi8t();
  vi8t table;
  table.reserve(n + 1);
  table.push_back(1);
  i8t r = 1;
  F2(i, 1, n + 1) table.push_back(r = r * i % mod);
  return table;
}
i8t mod_nCr(u8t n, u8t r, u8t mod, vi8t const &table) {
  i8t num = mod_factorial(n, mod);
  i8t d1 = mod_factorial(r, mod);
  i8t d2 = mod_factorial(n - r, mod);
  return mod_div(mod_div(num, d1, mod), d2, mod);
}

//-----------------------------------------------------------------------------

static constexpr i8t mod = 1e9 + 7; // これは素数である。

i8t calculate_combinat(vector<pair<i8t, i8t>> const &data, i8t mod,
                       vi8t const &fac_table) {
  // 多項係数
  i8t wt = 0;
  i8t den = 1;
  for (auto p : data) {
    i8t w = p.first;
    wt += w;
    den = den * fac_table[w] % mod;
  }
  u8t npat = mod_div(fac_table[wt], den, mod);

  // 子木の生成組み合わせ数を掛ける。
  for (auto p : data)
    npat = npat * p.second % mod;

  return npat;
}

int main() {
  R(N);
  std::vector<node_t> vtx(N);
  F(I, N - 1) {
    R(u1) R(v1);
    u1--;
    v1--;
    vtx[u1].adj.push_back(v1);
    vtx[v1].adj.push_back(u1);
  }
  tree_make_directed_by_root(vtx, 0);

  auto factab = mod_factorial_table(N, mod);
  auto _set_parent = [&](auto p, auto parent) { vtx[p].parent = parent; };
  auto _update_npat1 = [&](auto p, auto) {
    auto &node = vtx[p];
    vector<pair<i8t, i8t>> data;
    for (auto i : node.adj)
      data.emplace_back(vtx[i].weight, vtx[i].npat_child);
    node.npat_child = calculate_combinat(data, mod, factab);
  };

  auto _update_npat2 = [&](auto p, auto parent) {
    if (parent < 0)
      return;

    auto &node = vtx[parent];
    i8t w0 = node.weight;
    i8t w1 = vtx[p].weight;
    i8t n1 = vtx[p].npat_child;
    i8t nW = node.parent < 0 ? 1 : node.npat_parent;

    i8t npat = node.npat_child;
    npat = npat * factab[N - 1 - w1] % mod;
    npat = npat * factab[w1] % mod;
    npat = mod_div(npat, factab[w0 - 1], mod);
    npat = mod_div(npat, factab[N - w0], mod);
    npat = npat * nW % mod;
    npat = mod_div(npat, n1, mod);

    vtx[p].npat_parent = npat;
  };

  tree_visit(vtx, 0, _set_parent, _update_npat1);
  tree_visit(vtx, 0, _update_npat2, [](auto, auto) {});

  F(i, N) {
    auto &node = vtx[i];
    vector<pair<i8t, i8t>> data;
    for (auto i : node.adj)
      data.emplace_back(vtx[i].weight, vtx[i].npat_child);
    if (node.parent >= 0) {
      i8t w = vtx.size() - node.weight;
      data.emplace_back(w, node.npat_parent);
    }
    P(calculate_combinat(data, mod, factab));
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <numeric>
#include <string>
#include <utility>
#include <vector>
using namespace std;
using i8t = std::int64_t;
using u8t = std::uint64_t;
using vi8t = std::vector<i8t>;
using vvi8t = std::vector<std::vector<i8t>>;
using vu8t = std::vector<u8t>;
#include <limits>
static const i8t i8t_max = numeric_limits<i8t>::max();
static const i8t i8t_min = numeric_limits<i8t>::min();
#include <unordered_map>
#include <unordered_set>
template <typename Key, typename Value>
using map_t = std::unordered_map<Key, Value>;
template <typename Key> using set_t = std::unordered_set<Key>;
#define F3(i, i0, iN, step) for (i8t i = i0, i##_N = iN; i < i##_N; i += step)
#define F2(i, i0, iN) F3(i, i0, iN, 1)
#define F(i, N) F3(i, 0, N, 1)
#define FR(i, N) for (i8t i = N; i--;)
#define R(V)                                                                   \\
  i8t V;                                                                       \\
  std::cin >> V;
#define RV(V, N)                                                               \\
  std::vector<i8t> V;                                                          \\
  F(_i, N) { R(_x) V.push_back(_x); }
#define RL1(line) std::getline(std::cin, line);
#define RL(line)                                                               \\
  std::string line;                                                            \\
  RL1(line)
#define RM(V, L, C)                                                            \\
  std::vector<char> V;                                                         \\
  F(_i, L) { RL(_l) F(_j, C) V.push_back(_l[_j]); }
#define P(X) std::cout << (X) << std::endl;
template <typename T> T umax(T &ans, T const &value) {
  if (value > ans)
    ans = value;
  return ans;
}
template <typename T> T umin(T &ans, T const &value) {
  if (value < ans)
    ans = value;
  return ans;
}
template <typename T> struct identity {
  using type = T;
};

struct node_t {
  i8t weight = 0; // 子孫の数(自分も入れる)

  i8t parent = -1;
  std::vector<i8t> adj;

  i8t npat_child;
  i8t npat_parent;
};

void tree_make_directed_by_root(std::vector<node_t> &vtx, i8t root) {
  std::vector<bool> mark(vtx.size());
  std::vector<i8t> stk;
  std::vector<i8t> ind;
  i8t p, i;
  mark[p = root] = 1;
  i = 0;
  while (1) {
    node_t &node = vtx[p];
    if (i < node.adj.size()) {
      i8t &p2 = node.adj[i];
      if (mark[p2]) {
        p2 = -1;
        i++;
      } else {
        stk.push_back(p);
        ind.push_back(i + 1);
        mark[p = p2] = 1;
        i = 0;
      }
    } else {
      // remove parents
      node.adj.erase(std::remove(node.adj.begin(), node.adj.end(), -1),
                     node.adj.end());

      // update weight
      node.weight = 1;
      for (i8t p2 : node.adj)
        node.weight += vtx[p2].weight;

      if (stk.size()) {
        p = stk.back();
        i = ind.back();
        stk.pop_back();
        ind.pop_back();
      } else {
        break;
      }
    }
  }
}

template <typename FEnter, typename FLeave>
void tree_visit(std::vector<node_t> const &vtx, i8t root, FEnter _enter,
                FLeave _leave) {
  std::vector<i8t> stk, ind;
  i8t p, i;
  p = root;
  i = 0;
  _enter(p, -1);
  while (1) {
    auto &node = vtx[p];
    if (i < node.adj.size()) {
      stk.push_back(p);
      ind.push_back(i + 1);
      _enter(node.adj[i], p);
      p = node.adj[i];
      i = 0;
    } else {
      if (stk.size()) {
        _leave(p, stk.back());
        p = stk.back();
        i = ind.back();
        stk.pop_back();
        ind.pop_back();
      } else {
        _leave(p, -1);
        break;
      }
    }
  }
}

u8t u8t_gcd(u8t a, u8t b) {
  if (a < b)
    swap(a, b);
  if (b == 0)
    return 0;
  for (;;) {
    a %= b;
    if (a == 0)
      return b;
    b %= a;
    if (b == 0)
      return a;
  }
}

i8t i8t_gcd(i8t a, i8t b, i8t &s, i8t &t) {
  if (a == 0 || b == 0) {
    s = t = 0;
    return 0;
  } else if (a < 0) {
    i8t r = i8t_gcd(-a, b, s, t);
    s = -s;
    return r;
  } else if (b < 0) {
    i8t r = i8t_gcd(a, -b, s, t);
    t = -t;
    return r;
  }

  i8t q, as = 1, at = 0, bs = 0, bt = 1;
  for (;;) {
    q = a / b;
    a %= b;
    if (a == 0) {
      s = bs;
      t = bt;
      return b;
    }
    as -= q * bs;
    at -= q * bt;

    q = b / a;
    b %= a;
    if (b == 0) {
      s = as;
      t = at;
      return a;
    }
    bs -= q * as;
    bt -= q * at;
  }
}

/// @fn mod_inv(x,m)
///   互いに素な整数 x,m に対して x^{-1} mod m を求める。
/// @param[in] x
/// @param[in] m
i8t mod_inv(i8t x, i8t mod) {
  i8t xinv, b;
  if (1 != i8t_gcd(x, mod, xinv, b))
    std::exit(2); // 互いに素でない。
  xinv %= mod;
  if (xinv < 0)
    xinv += mod;
  return xinv;
}

i8t mod_div(i8t x, i8t y, i8t mod) { return x * mod_inv(y, mod) % mod; }

i8t mod_factorial(i8t n, i8t mod) {
  if (n <= 0)
    return n < 0 ? 0 : 1;
  i8t r = n;
  while (--n)
    r = r * n % mod;
  return r;
}

i8t mod_nCr(u8t n, u8t r, u8t mod) {
  i8t num = mod_factorial(n, mod);
  i8t d1 = mod_factorial(r, mod);
  i8t d2 = mod_factorial(n - r, mod);
  return mod_div(mod_div(num, d1, mod), d2, mod);
}

static constexpr i8t mod = 1e9 + 7; // これは素数である。

std::vector<std::pair<u8t, u8t>> prime_multifactors(u8t value) {
  std::vector<std::pair<u8t, u8t>> factors;
  if (value == 0) {
    factors.emplace_back(0, 1);
    return factors;
  }
  auto _reduce = [&](auto f) {
    if (value % f)
      return;
    value /= f;
    factors.emplace_back(f, 1);
    while (value % f == 0) {
      value /= f;
      factors.back().second++;
    }
  };
  _reduce(2);
  _reduce(3);
  for (i8t i = 5; i * i <= value; i += 6) {
    if (i * i > value)
      break;
    _reduce(i);
    _reduce(i + 2);
  }
  if (value != 1)
    factors.emplace_back(value, 1);
  return factors;
}
void debug1() {
  // for(auto p:prime_multifactors(mod)){
  //   cerr<<" "<<p.first;
  //   if(p.second>1)cerr<<"^"<<p.second;
  // }

  // Test gcd
  // i8t a,b;
  // i8t_gcd(5,11,a,b);
  // cerr<<a<<" "<<b<<endl;

  // Test factorial
  // cerr<<mod_factorial(5,mod)<<endl;

  // Test factorial
  // cerr<<mod_nCr(4,2,mod)<<endl;
}

//-----------------------------------------------------------------------------

i8t calculate_combinat(vector<pair<i8t, i8t>> const &data, i8t mod) {
  // 多項係数
  i8t wt = 0;
  i8t den = 1;
  for (auto p : data) {
    i8t w = p.first;
    wt += w;
    den = den * mod_factorial(w, mod) % mod;
  }
  u8t npat = mod_div(mod_factorial(wt, mod), den, mod);

  // 子木の生成組み合わせ数を掛ける。
  for (auto p : data)
    npat = npat * p.second % mod;

  return npat;
}

int main() {
  // debug1();return 1;

  R(N);
  std::vector<node_t> vtx(N);
  F(I, N - 1) {
    R(u1) R(v1);
    u1--;
    v1--;
    vtx[u1].adj.push_back(v1);
    vtx[v1].adj.push_back(u1);
  }
  tree_make_directed_by_root(vtx, 0);

  auto _set_parent = [&](auto p, auto parent) { vtx[p].parent = parent; };
  auto _update_npat1 = [&](auto p, auto) {
    auto &node = vtx[p];
    vector<pair<i8t, i8t>> data;
    for (auto i : node.adj)
      data.emplace_back(vtx[i].weight, vtx[i].npat_child);
    node.npat_child = calculate_combinat(data, mod);
  };

  auto _update_npat2 = [&](auto p, auto parent) {
    if (parent < 0)
      return;

    auto &node = vtx[parent];

    vector<pair<i8t, i8t>> data;
    for (auto i : node.adj) {
      if (i == p)
        continue;
      data.emplace_back(vtx[i].weight, vtx[i].npat_child);
    }
    if (node.parent >= 0) {
      i8t w = vtx.size() - node.weight;
      data.emplace_back(w, node.npat_parent);
    }
    vtx[p].npat_parent = calculate_combinat(data, mod);
  };

  tree_visit(vtx, 0, _set_parent, _update_npat1);
  tree_visit(vtx, 0, _update_npat2, [](auto, auto) {});

  F(i, N) {
    auto &node = vtx[i];
    vector<pair<i8t, i8t>> data;
    for (auto i : node.adj)
      data.emplace_back(vtx[i].weight, vtx[i].npat_child);
    if (node.parent >= 0) {
      i8t w = vtx.size() - node.weight;
      data.emplace_back(w, node.npat_parent);
    }
    P(calculate_combinat(data, mod));
  }
  return 0;
}