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

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

#define int long long
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, n) for (int i = 1; i <= (n); i++)
#define all(x) begin(x), end(x)
#define uniq(x) (x).erase(unique(all(x)), end(x))
#define bit(n) (1LL << (n))
#define cdiv(a, b) (((a)-1) / (b) + 1)
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr long double PI = 3.1415926535897932384626433832795028L;
constexpr int DY[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int DX[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T> void fin(T mes) {
  cout << mes << endl;
  exit(0);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U> bool chmin(T &a, const U &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

struct re_rooting {
  int n;
  std::vector<std::vector<int>> graph;
  re_rooting(int n) : n(n), graph(n) {}
  void add_edge(int u, int v) { graph[u].push_back(v), graph[v].push_back(u); }
  template <typename T, typename F, typename G>
  std::vector<T> get(T identity, F f, G g) {
    std::vector<std::vector<T>> dp(n, std::vector<T>());
    {
      auto dfs = [&](auto &&self, int cur, int par) -> T {
        dp[cur].resize(graph[cur].size());
        T cum = identity;
        for (int i = 0; i < graph[cur].size(); i++) {
          if (graph[cur][i] == par) {
            continue;
          }
          dp[cur][i] = g(self(self, graph[cur][i], cur), cur, i);
          cum = f(cum, dp[cur][i]);
        }
        return cum;
      };
      dfs(dfs, 0, -1);
    }
    std::vector<T> ret(n);
    {
      auto dfs = [&](auto &&self, int cur, int par, T drop) -> void {
        std::vector<T> l_cum(graph[cur].size() + 1),
            r_cum(graph[cur].size() + 1);
        l_cum[0] = r_cum[graph[cur].size()] = identity;
        for (int i = 0; i < graph[cur].size(); i++) {
          if (graph[cur][i] == par) {
            dp[cur][i] = g(drop, cur, i);
          }
          l_cum[i + 1] = f(l_cum[i], dp[cur][i]);
        }
        for (int i = graph[cur].size() - 1; i >= 0; i--) {
          r_cum[i] = f(dp[cur][i], r_cum[i + 1]);
        }
        for (int i = 0; i < graph[cur].size(); i++) {
          if (graph[cur][i] == par) {
            continue;
          }
          self(self, graph[cur][i], cur, f(l_cum[i], r_cum[i + 1]));
        }
        ret[cur] = l_cum[graph[cur].size()];
      };
      dfs(dfs, 0, -1, identity);
    }
    return ret;
  }
};

template <int MOD = 1000000007> struct mod_int {
  int val;
  mod_int(long long val_ = 0) : val(val_ % MOD) {
    if (val < 0) {
      val += MOD;
    }
  }
  bool operator==(const mod_int &rhs) const { return val == rhs.val; }
  bool operator!=(const mod_int &rhs) const {
    return std::rel_ops::operator!=(*this, rhs);
  }
  mod_int &operator+=(const mod_int &rhs) {
    if ((val += rhs.val) >= MOD) {
      val -= MOD;
    }
    return *this;
  }
  mod_int &operator-=(const mod_int &rhs) {
    if ((val += MOD - rhs.val) >= MOD) {
      val -= MOD;
    }
    return *this;
  }
  mod_int &operator*=(const mod_int &rhs) {
    val = (int)((long long)val * rhs.val % MOD);
    return *this;
  }
  mod_int &operator/=(const mod_int &rhs) { return *this *= rhs.inv(); }
  mod_int operator+() const { return *this; }
  mod_int operator-() const { return -val; }
  mod_int operator++() { return *this += 1; }
  mod_int operator--() { return *this -= 1; }
  mod_int operator++(signed) {
    const mod_int ret(*this);
    ++*this;
    return ret;
  }
  mod_int operator--(signed) {
    const mod_int ret(*this);
    --*this;
    return ret;
  }
  mod_int operator+(const mod_int &rhs) const { return mod_int(*this) += rhs; }
  mod_int operator-(const mod_int &rhs) const { return mod_int(*this) -= rhs; }
  mod_int operator*(const mod_int &rhs) const { return mod_int(*this) *= rhs; }
  mod_int operator/(const mod_int &rhs) const { return mod_int(*this) /= rhs; }
  mod_int inv() const {
    int a = val, b = MOD, x = 1, u = 0;
    while (b) {
      int t = a / b;
      std::swap(a -= t * b, b);
      std::swap(x -= t * u, u);
    }
    return x;
  }
  mod_int pow(long long n) const {
    if (n < 0) {
      return pow(-n).inv();
    }
    mod_int ret = 1, mul = *this;
    while (n) {
      if (n & 1) {
        ret *= mul;
      }
      mul *= mul;
      n >>= 1;
    }
    return ret;
  }
  friend std::istream &operator>>(std::istream &is, mod_int &rhs) {
    long long v;
    is >> v;
    rhs = v;
    return is;
  }
  friend std::ostream &operator<<(std::ostream &os, const mod_int &rhs) {
    return os << rhs.val;
  }
  struct combination {
    std::vector<mod_int> fact{1, 1}, f_inv{1, 1}, inv{0, 1};
    void calc(int n) {
      while (fact.size() <= n) {
        int i = fact.size();
        fact.push_back(fact[i - 1] * i);
        inv.push_back(-inv[MOD % i] * (MOD / i));
        f_inv.push_back(f_inv[i - 1] * inv[i]);
      }
    }
    mod_int P(int n, int r) {
      return r < 0 || n < r ? 0 : (calc(n), fact[n] * f_inv[n - r]);
    }
    mod_int C(int n, int r) {
      return r < 0 || n < r ? 0 : (calc(n), fact[n] * f_inv[r] * f_inv[n - r]);
    }
    mod_int H(int n, int r) { return C(n + r - 1, r); }
  };
};

using mint = mod_int<>;
mint::combination comb;
int N;
int a[220000], b[220000];
signed main() {
  cin >> N;
  re_rooting rr(N);
  rep(i, N - 1) {
    cin >> a[i] >> b[i];
    a[i]--, b[i]--;
    rr.add_edge(a[i], b[i]);
  }
  using wombat = pair<mint, int>;
  vector<wombat> ans = rr.get(
      wombat(1, 0),
      [](wombat a, wombat b) {
        return make_pair(a.first * b.first *
                             comb.C(a.second + b.second, a.second),
                         a.second + b.second);
      },
      [](wombat a, int, int) { return wombat(a.first, a.second + 1); });
  rep(i, N) { cout << ans[i].first << endl; }
}


Error: TLE

Fixed submission:

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

#define int long long
#define double long double
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, n) for (int i = 1; i <= (n); i++)
#define all(x) begin(x), end(x)
#define uniq(x) (x).erase(unique(all(x)), end(x))
#define bit(n) (1LL << (n))
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr double pi = 3.1415926535897932384626433832795028L;
constexpr int dy[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int dx[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T> void fin(T mes) {
  cout << mes << endl;
  exit(0);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U> bool chmin(T &a, const U &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

int N;
int a[220000], b[220000];
vint e[220000];
map<int, int> dp1[220000];
map<int, int> dp2[220000];
int n[220000];
int ans[220000];
int mod = 1e9 + 7;
int fact[220000], f_inv[220000], inv[220000];
int C(int n, int r) { return fact[n] * f_inv[r] % mod * f_inv[n - r] % mod; }
pair<int, int> dfs(int x, int p) {
  if (dp1[x].count(p) != 0) {
    return {dp1[x][p], dp2[x][p]};
  }
  int ret1 = 1;
  int ret2 = 0;
  int c;
  if (p == a[x]) {
    c = b[x];
  } else {
    c = a[x];
  }
  rep(i, e[c].size()) {
    if (e[c][i] == x) {
      continue;
    }
    pair<int, int> r = dfs(e[c][i], c);
    (ret1 *= r.first) %= mod;
    (ret1 *= C(ret2 + r.second, ret2)) %= mod;
    ret2 += r.second;
  }
  ret2++;
  (ans[p] *= ret1) %= mod;
  (ans[p] *= C(n[p] + ret2, n[p])) %= mod;
  n[p] += ret2;
  return {dp1[x][p] = ret1, dp2[x][p] = ret2};
}
signed main() {
  cin >> N;
  rep(i, N - 1) {
    cin >> a[i] >> b[i];
    a[i]--, b[i]--;
    e[a[i]].push_back(i), e[b[i]].push_back(i);
  }
  rep(i, N) { ans[i] = 1; }

  fact[0] = fact[1] = 1;
  inv[1] = 1;
  f_inv[0] = f_inv[1] = 1;
  for (int i = 2; i < 220000; i++) {
    fact[i] = fact[i - 1] * i % mod;
    inv[i] = (-inv[mod % i] + mod) * (mod / i) % mod;
    f_inv[i] = f_inv[i - 1] * inv[i] % mod;
  }

  rep(i, N - 1) {
    dfs(i, a[i]);
    dfs(i, b[i]);
  }
  rep(i, N) { cout << ans[i] << endl; }
}