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

//{{{
#include <algorithm>
#include <cmath>
#include <deque>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <sys/time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using ll = long long;
enum : int { M = (int)1e9 + 7 };
enum : ll { MLL = (ll)1e18L + 9 };
using namespace std;
#ifdef LOCAL
#include "rprint2.hpp"
#else
#define FUNC(name)                                                             \\
  template <ostream &out = cout, class... T> void name(T &&...) {}
FUNC(printde) FUNC(printdbe) FUNC(printdwe) FUNC(printdu);
#endif
template <template <class T, class = std::allocator<T>> class V, class E>
istream &operator>>(istream &in, V<E> &v) {
  for (auto &e : v) {
    in >> e;
  }
  return in;
}
//}}}

// mint
template <int M> struct ModInt {
  int num;
  ModInt() : num(0) {}
  ModInt(int num) : num(num) {}
  operator int() const { return num; }
  // template <class T> ModInt operator + (T n) const { n += num; return n >= M
  // ? n - M : n; } template <class T> ModInt operator - (T n) const { n = num -
  // n; return n < 0 ? n + M : n; }
  template <class T> ModInt operator+(T n) const {
    auto ret = num + n;
    return ret >= M ? ret - M : ret;
  }
  template <class T> ModInt operator-(T n) const {
    auto ret = num - n;
    return ret < 0 ? ret + M : ret;
  }
  // template <class T> ModInt operator + (T n) const { return (num + n) % M; }
  // template <class T> ModInt operator - (T n) const { return (num - n + M) %
  // M; }
  template <class T> ModInt operator*(T n) const { return (ll)num * n % M; }
  template <class T> ModInt operator/(T n) const {
    return (ll)num * modInv(ModInt(n)) % M;
  }
  template <class T> ModInt operator>>(T n) const { return num >> n; }
  template <class T> ModInt operator<<(T n) const { return ((ll)num << n) % M; }
  template <class T> ModInt operator|(T n) const { return num | n; }
  template <class T> ModInt operator&(T n) const { return num & n; }
  template <class T> bool operator<(T n) const { return num < n; }
  template <class T> bool operator>(T n) const { return num > n; }
  template <class T> bool operator<=(T n) const { return num <= n; }
  template <class T> bool operator>=(T n) const { return num >= n; }
  template <class T> bool operator==(T n) const { return num == n; }
  template <class T> bool operator!=(T n) const { return num != n; }
  ModInt operator++(int) {
    int ret = num;
    num = num == M + 1 ? 0 : num + 1;
    return ret;
  }
  ModInt operator--(int) {
    int ret = num;
    num = num == 0 ? M - 1 : num - 1;
    return ret;
  }
  ModInt operator++() { return num = num == M + 1 ? 0 : num + 1; }
  ModInt operator--() { return num = num == 0 ? M - 1 : num - 1; }
  // ModInt operator ++ (int){ int ret = num; num = (num + 1) % M; return ret; }
  // ModInt operator -- (int){ int ret = num; num = (num - 1 + M) % M; return
  // ret; } ModInt operator ++ (){ return num = (num + 1) % M; } ModInt operator
  // -- (){ return num = (num - 1 + M) % M; }
  template <class T> ModInt operator+=(T n) {
    num += n;
    return num = num >= M ? num - M : num;
  }
  template <class T> ModInt operator-=(T n) {
    num -= n;
    return num = num < 0 ? num + M : num;
  }
  // template <class T> ModInt operator += (T n){ auto ret = num + n; return num
  // = ret >= M ? ret - M : ret; } template <class T> ModInt operator -= (T n){
  // auto ret = num - n; return num = ret < 0 ? ret + M : ret; } template <class
  // T> ModInt operator += (T n){ return num = (num + n) % M; } template <class
  // T> ModInt operator -= (T n){ return num = (num - n) % M; }
  template <class T> ModInt operator*=(T n) { return num = (ll)num * n % M; }
  template <class T> ModInt operator/=(T n) {
    return num = (ll)num * modInv(n) % M;
  }
  template <class T> ModInt operator>>=(T n) { return num = num >> n; }
  template <class T> ModInt operator<<=(T n) {
    return num = ((ll)num << n) % M;
  }
  friend ModInt modPow(ModInt base, ModInt ex);
  static vector<ModInt> mFacts, mInvFacts;
  friend ModInt modFact(ModInt n) { return mFacts[n.num]; }
  friend ModInt modInvFact(ModInt n) { return mInvFacts[n.num]; }
  friend ModInt modInv(ModInt n) { return modPow(n, M - 2); }
  friend ModInt modC(ModInt n, ModInt k) {
    return modFact(n) * modInvFact(k) * modInvFact(n - k);
  }
  static void init(int n) {
    mFacts.resize(n + 1);
    mInvFacts.resize(n + 1);
    mFacts[0] = 1;
    for (int i = 1; i <= n; i++) {
      mFacts[i] = mFacts[i - 1] * i;
    }
    mInvFacts[n] = modInv(mFacts[n]);
    for (int i = n - 1; i >= 0; i--) {
      mInvFacts[i] = mInvFacts[i + 1] * (i + 1);
    }
  }
  friend istream &operator>>(istream &in, ModInt &e) { return in >> e.num; }
  friend ostream &operator<<(ostream &out, const ModInt &e) {
    return out << e.num;
  }
};
template <int M> vector<ModInt<M>> ModInt<M>::mFacts;
template <int M> vector<ModInt<M>> ModInt<M>::mInvFacts;
// using mint = ModInt<998244353>;
using mint = ModInt<(int)M>;
mint modPow(mint base, mint ex) {
  mint ret = 1;
  while (ex.num) {
    if (ex.num & 1) {
      ret = ret * base;
    }
    ex >>= 1;
    base = base * base;
  }
  return ret;
}

struct Graph {
  struct Node {
    vector<int> dests;
    unordered_map<int, pair<mint, int>> valueNums;
    // unordered_map<int, mint> values;
    // unordered_map<int, int> nums;
  };
  vector<Node> nodes;
  Graph(int num) : nodes(num) {}
  int id(Node *n) { return n - &nodes[0]; }
  int id(Node &n) { return &n - &nodes[0]; }
  // void print(){ for(auto& node : nodes){ prints(id(node), node.dests); } }
  void add1Edge(int from, int dest) { nodes[from].dests.push_back(dest); }
  void add2Edge2(int from, int dest) {
    add1Edge(from, dest);
    add1Edge(dest, from);
  }
  void dfs(int idx, int from = -1) {
    if (from != -1 && nodes[idx].valueNums.count(-1)) {
      auto p1 = nodes[idx].valueNums[-1];
      mint acc = p1.first;
      int num = p1.second;
      num--;
      acc *= modFact((mint)nodes[from].valueNums[idx].second);
      acc /= modFact((mint)num);
      // auto p = nodes[idx].valueNums[from];
      auto p = nodes[from].valueNums[idx];
      acc /= p.first;
      num -= p.second;
      acc *= modFact((mint)num);
      num++;
      nodes[idx].valueNums[from] = pair<mint, int>(acc, num);
      return;
    }
    mint acc = 1;
    int num = 0;
    for (int didx : nodes[idx].dests) {
      if (didx == from) {
        continue;
      }
      if (!nodes[didx].valueNums.count(idx)) {
        dfs(didx, idx);
      }
      auto p = nodes[didx].valueNums[idx];
      acc *= p.first;
      num += p.second;
    }
    acc *= modFact((mint)num);
    for (int didx : nodes[idx].dests) {
      if (didx == from) {
        continue;
      }
      acc /= modFact((mint)nodes[didx].valueNums[idx].second);
    }
    num++;
    nodes[idx].valueNums[from] = pair<mint, int>(acc, num);
  }
  // void dfs(int idx, int from = -1){
  //     mint acc = 1;
  //     int num = 0;
  //     for(int didx : nodes[idx].dests){
  //         if(didx == from){ continue; }
  //         if(!nodes[didx].valueNums.count(idx)){
  //             dfs(didx, idx);
  //         }
  //         auto p = nodes[didx].valueNums[idx];
  //         acc *= p.first;
  //         num += p.second;
  //         // acc *= nodes[didx].values[idx];
  //         // num += nodes[didx].nums[idx];
  //     }
  //     // acc *= modFact((mint)(nodes[idx].dests.size() - (from != -1)));
  //     acc *= modFact((mint)num);
  //     for(int didx : nodes[idx].dests){
  //         if(didx == from){ continue; }
  //         acc /= modFact((mint)nodes[didx].valueNums[idx].second);
  //         // acc /= modFact((mint)nodes[didx].nums[idx]);
  //     }
  //     num++;
  //     nodes[idx].valueNums[from] = pair<mint, int>(acc, num);
  //     // nodes[idx].values[from] = acc;
  //     // nodes[idx].nums[from] = num;
  // }
};

int main() {
  mint::init(3e5);
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  Graph g(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g.add2Edge2(a, b);
  }
  vector<int> is(n);
  iota(is.begin(), is.end(), 0);
  sort(is.begin(), is.end(), [&](int i, int j) {
    return g.nodes[i].dests.size() > g.nodes[j].dests.size();
  });
  for (auto &i : is) {
    g.dfs(i);
  }
  for (int i = 0; i < n; i++) {
    // if(!g.nodes[i].valueNums.count(-1)){
    //     g.dfs(i);
    // }
    // g.dfs(i);
    cout << g.nodes[i].valueNums[-1].first << '\\n';
    // cout << g.nodes[i].values[-1] << '\\n';
  }
}


Error: TLE

Fixed submission:

//{{{
#include <algorithm>
#include <cmath>
#include <deque>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <sys/time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using ll = long long;
enum : int { M = (int)1e9 + 7 };
enum : ll { MLL = (ll)1e18L + 9 };
using namespace std;
#ifdef LOCAL
#include "rprint2.hpp"
#else
#define FUNC(name)                                                             \\
  template <ostream &out = cout, class... T> void name(T &&...) {}
FUNC(printde) FUNC(printdbe) FUNC(printdwe) FUNC(printdu);
#endif
template <template <class T, class = std::allocator<T>> class V, class E>
istream &operator>>(istream &in, V<E> &v) {
  for (auto &e : v) {
    in >> e;
  }
  return in;
}
//}}}

// mint
template <int M> struct ModInt {
  int num;
  ModInt() : num(0) {}
  ModInt(int num) : num(num) {}
  operator int() const { return num; }
  // template <class T> ModInt operator + (T n) const { n += num; return n >= M
  // ? n - M : n; } template <class T> ModInt operator - (T n) const { n = num -
  // n; return n < 0 ? n + M : n; }
  template <class T> ModInt operator+(T n) const {
    auto ret = num + n;
    return ret >= M ? ret - M : ret;
  }
  template <class T> ModInt operator-(T n) const {
    auto ret = num - n;
    return ret < 0 ? ret + M : ret;
  }
  // template <class T> ModInt operator + (T n) const { return (num + n) % M; }
  // template <class T> ModInt operator - (T n) const { return (num - n + M) %
  // M; }
  template <class T> ModInt operator*(T n) const { return (ll)num * n % M; }
  template <class T> ModInt operator/(T n) const {
    return (ll)num * modInv(ModInt(n)) % M;
  }
  template <class T> ModInt operator>>(T n) const { return num >> n; }
  template <class T> ModInt operator<<(T n) const { return ((ll)num << n) % M; }
  template <class T> ModInt operator|(T n) const { return num | n; }
  template <class T> ModInt operator&(T n) const { return num & n; }
  template <class T> bool operator<(T n) const { return num < n; }
  template <class T> bool operator>(T n) const { return num > n; }
  template <class T> bool operator<=(T n) const { return num <= n; }
  template <class T> bool operator>=(T n) const { return num >= n; }
  template <class T> bool operator==(T n) const { return num == n; }
  template <class T> bool operator!=(T n) const { return num != n; }
  ModInt operator++(int) {
    int ret = num;
    num = num == M + 1 ? 0 : num + 1;
    return ret;
  }
  ModInt operator--(int) {
    int ret = num;
    num = num == 0 ? M - 1 : num - 1;
    return ret;
  }
  ModInt operator++() { return num = num == M + 1 ? 0 : num + 1; }
  ModInt operator--() { return num = num == 0 ? M - 1 : num - 1; }
  // ModInt operator ++ (int){ int ret = num; num = (num + 1) % M; return ret; }
  // ModInt operator -- (int){ int ret = num; num = (num - 1 + M) % M; return
  // ret; } ModInt operator ++ (){ return num = (num + 1) % M; } ModInt operator
  // -- (){ return num = (num - 1 + M) % M; }
  template <class T> ModInt operator+=(T n) {
    num += n;
    return num = num >= M ? num - M : num;
  }
  template <class T> ModInt operator-=(T n) {
    num -= n;
    return num = num < 0 ? num + M : num;
  }
  // template <class T> ModInt operator += (T n){ auto ret = num + n; return num
  // = ret >= M ? ret - M : ret; } template <class T> ModInt operator -= (T n){
  // auto ret = num - n; return num = ret < 0 ? ret + M : ret; } template <class
  // T> ModInt operator += (T n){ return num = (num + n) % M; } template <class
  // T> ModInt operator -= (T n){ return num = (num - n) % M; }
  template <class T> ModInt operator*=(T n) { return num = (ll)num * n % M; }
  template <class T> ModInt operator/=(T n) {
    return num = (ll)num * modInv(n) % M;
  }
  template <class T> ModInt operator>>=(T n) { return num = num >> n; }
  template <class T> ModInt operator<<=(T n) {
    return num = ((ll)num << n) % M;
  }
  friend ModInt modPow(ModInt base, ModInt ex);
  static vector<ModInt> mFacts, mInvFacts;
  friend ModInt modFact(ModInt n) { return mFacts[n.num]; }
  friend ModInt modInvFact(ModInt n) { return mInvFacts[n.num]; }
  friend ModInt modInv(ModInt n) { return modPow(n, M - 2); }
  friend ModInt modC(ModInt n, ModInt k) {
    return modFact(n) * modInvFact(k) * modInvFact(n - k);
  }
  static void init(int n) {
    mFacts.resize(n + 1);
    mInvFacts.resize(n + 1);
    mFacts[0] = 1;
    for (int i = 1; i <= n; i++) {
      mFacts[i] = mFacts[i - 1] * i;
    }
    mInvFacts[n] = modInv(mFacts[n]);
    for (int i = n - 1; i >= 0; i--) {
      mInvFacts[i] = mInvFacts[i + 1] * (i + 1);
    }
  }
  friend istream &operator>>(istream &in, ModInt &e) { return in >> e.num; }
  friend ostream &operator<<(ostream &out, const ModInt &e) {
    return out << e.num;
  }
};
template <int M> vector<ModInt<M>> ModInt<M>::mFacts;
template <int M> vector<ModInt<M>> ModInt<M>::mInvFacts;
// using mint = ModInt<998244353>;
using mint = ModInt<(int)M>;
mint modPow(mint base, mint ex) {
  mint ret = 1;
  while (ex.num) {
    if (ex.num & 1) {
      ret = ret * base;
    }
    ex >>= 1;
    base = base * base;
  }
  return ret;
}

struct Graph {
  struct Node {
    vector<int> dests;
    unordered_map<int, pair<mint, int>> valueNums;
    // unordered_map<int, mint> values;
    // unordered_map<int, int> nums;
  };
  vector<Node> nodes;
  Graph(int num) : nodes(num) {}
  int id(Node *n) { return n - &nodes[0]; }
  int id(Node &n) { return &n - &nodes[0]; }
  // void print(){ for(auto& node : nodes){ prints(id(node), node.dests); } }
  void add1Edge(int from, int dest) { nodes[from].dests.push_back(dest); }
  void add2Edge2(int from, int dest) {
    add1Edge(from, dest);
    add1Edge(dest, from);
  }
  void dfs(int idx, int from = -1) {
    mint acc = 1;
    int num = 0;
    for (int didx : nodes[idx].dests) {
      if (didx == from) {
        continue;
      }
      if (!nodes[didx].valueNums.count(idx)) {
        dfs(didx, idx);
      }
      auto p = nodes[didx].valueNums[idx];
      acc *= p.first;
      num += p.second;
      // acc *= nodes[didx].values[idx];
      // num += nodes[didx].nums[idx];
    }
    // acc *= modFact((mint)(nodes[idx].dests.size() - (from != -1)));
    acc *= modFact((mint)num);
    for (int didx : nodes[idx].dests) {
      if (didx == from) {
        continue;
      }
      acc /= modFact((mint)nodes[didx].valueNums[idx].second);
      // acc /= modFact((mint)nodes[didx].nums[idx]);
    }
    num++;
    nodes[idx].valueNums[from] = pair<mint, int>(acc, num);
    // nodes[idx].values[from] = acc;
    // nodes[idx].nums[from] = num;
  }
};

int main() {
  mint::init(3e5);
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  Graph g(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    g.add2Edge2(a, b);
  }
  for (int i = 0; i < n; i++) {
    if (!g.nodes[i].valueNums.count(-1)) {
      g.dfs(i);
    }
    // g.dfs(i);
    cout << g.nodes[i].valueNums[-1].first << '\\n';
    // cout << g.nodes[i].values[-1] << '\\n';
  }
}