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

#include <iostream>
// #include <iomanip>
// #include <string>
#include <algorithm>
#include <utility>
#include <vector>
// #include <set>
#include <map>
// #include <queue>
// #include <deque>
// #include <bitset>
// #include <math.h>
using namespace std;
using ll = long long;
// using ld = long double ;
using vll = vector<ll>;
using vvll = vector<vll>;
// using vc = vecto\\r<char> ;
// using vvc = vector<vc> ;
// using vb = vector<bool> ;
// using vvb = vector<vb> ;
using pll = pair<ll, ll>;
// #define all(v) v.begin(),v.end()
ll mod = 1e9 + 7;
// long double pie = acos(-1) ;
// ll INF = 1e15 ;

// void yorn(bool a){if(a) cout << "Yes" << endl ; else cout << "No" << endl ;}
// void YorN(bool a){if(a) cout << "YES" << endl ; else cout << "NO" << endl ;}
// ll gcd(long long a,long long b){if(b==0) return a ; return gcd(b,a%b) ;}
// ll lcm(long long a,long long b){return a/gcd(a,b)*b ;}
// ll sa(long long a,long long b){if(a>b) return a-b ; return b-a ;}
// void fix_cout(){cout << fixed << setprecision(20) ;}

vll fac(2, 1), inv(2, 1), finv(2, 1);

void nCr_set(ll n) {
  for (int i = 2; i <= n; i++) {
    fac.emplace_back(fac.at(fac.size() - 1) * i % mod);
    inv.emplace_back(mod - inv.at(mod % i) * (mod / i) % mod);
    finv.emplace_back(finv.at(finv.size() - 1) * inv.at(i) % mod);
  }
}

ll nCr(ll n, ll r) {
  if (n < 0 || r < 0 || n < r)
    return 0;
  if (n == r || r == 0)
    return 1;
  return fac.at(n) * finv.at(n - r) % mod * finv.at(r) % mod;
}

ll nHr(ll n, ll r) { return nCr(n + r - 1, n - 1); }

ll modpow(ll a, ll b) {
  ll res = 1;
  while (b > 0) {
    if (b & 1)
      res = res * a % mod;
    a = a * a % mod;
    b >>= 1;
  }
  return res;
}

void DFS1(vvll &g, vll &sz, vll &par, int cur, int pre) {
  if (sz.at(cur) != -1)
    return;
  ll sum = 0;
  for (int i = 0; i < g.at(cur).size(); i++) {
    if (pre != g.at(cur).at(i)) {
      DFS1(g, sz, par, g.at(cur).at(i), cur);
      sum += sz.at(g.at(cur).at(i));
      par.at(g.at(cur).at(i)) = cur;
    }
  }
  sz.at(cur) = sum + 1;
}

void DFS2(vvll &g, vll &cnt, vll &sz, int cur) {
  if (cnt.at(cur) != -1)
    return;
  ll sum = 1;
  ll sub = 0;
  for (int i = 0; i < g.at(cur).size(); i++) {
    if (sz.at(cur) > sz.at(g.at(cur).at(i))) {
      DFS2(g, cnt, sz, g.at(cur).at(i));
      sum = sum * nCr(sz.at(cur) - 1 - sub, sz.at(g.at(cur).at(i))) % mod *
            cnt.at(g.at(cur).at(i)) % mod;
      sub += sz.at(g.at(cur).at(i));
    }
  }
  cnt.at(cur) = sum;
}

void DFS3(vvll &g, vll &rev, vll &cnt, vll &sz, vll &par, int cur) {
  if (rev.at(cur) != -1)
    return;
  ll sum = sz.at(par.at(cur));
  ll num = cnt.at(par.at(cur));
  num = num * finv.at(sum - 1) % mod;
  // cout << num << endl ;
  num = num * fac.at(sz.at(cur)) % mod;
  // cout << num << endl ;
  num = num * modpow(cnt.at(cur), mod - 2) % mod;
  // cout << num << endl ;
  ll add = g.size() - sz.at(par.at(cur));
  sum = sum - sz.at(cur) + add - 1 + mod;
  sum %= mod;
  num = num * fac.at(sum) % mod;
  num = num * finv.at(add) % mod;
  DFS3(g, rev, cnt, sz, par, par.at(cur));
  num = num * rev.at(par.at(cur)) % mod;
  rev.at(cur) = num;
}

int main() {
  ll n;
  cin >> n;
  nCr_set(2 * n);
  vvll g(n);
  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    g.at(a).push_back(b);
    g.at(b).push_back(a);
  }
  vll sz(n, -1), par(n, 0);
  DFS1(g, sz, par, 0, -1);
  vll cnt(n, -1), rev(n, -1);
  DFS2(g, cnt, sz, 0);
  rev.at(0) = 1;
  for (int i = 0; i < n; i++)
    DFS3(g, rev, cnt, sz, par, i);
  for (int i = 0; i < n; i++) {
    ll ans = 1;
    ll sum = 0;
    for (int j = 0; j < g.at(i).size(); j++) {
      if (par.at(g.at(i).at(j)) == i) {
        sum += sz.at(g.at(i).at(j));
        ans = ans * finv.at(sz.at(g.at(i).at(j))) % mod;
        ans = ans * cnt.at(g.at(i).at(j)) % mod;
      } else {
        sum += n - sz.at(i);
        ans = ans * finv.at(n - sz.at(i)) % mod;
        ans = ans * rev.at(i) % mod;
      }
    }
    ans = ans * fac.at(sum) % mod;
    cout << ans << endl;
  }
  // for(auto i:cnt) cout << i << " " ; cout << endl ;
  // for(auto i:rev) cout << i << " " ; cout << endl ;
}


Error: TLE

Fixed submission:

#include <iostream>
// #include <iomanip>
// #include <string>
#include <algorithm>
#include <utility>
#include <vector>
// #include <set>
#include <map>
// #include <queue>
// #include <deque>
// #include <bitset>
// #include <math.h>
using namespace std;
using ll = long long;
// using ld = long double ;
using vll = vector<ll>;
using vvll = vector<vll>;
// using vc = vecto\\r<char> ;
// using vvc = vector<vc> ;
// using vb = vector<bool> ;
// using vvb = vector<vb> ;
using pll = pair<ll, ll>;
// #define all(v) v.begin(),v.end()
ll mod = 1e9 + 7;
// long double pie = acos(-1) ;
// ll INF = 1e15 ;

// void yorn(bool a){if(a) cout << "Yes" << endl ; else cout << "No" << endl ;}
// void YorN(bool a){if(a) cout << "YES" << endl ; else cout << "NO" << endl ;}
// ll gcd(long long a,long long b){if(b==0) return a ; return gcd(b,a%b) ;}
// ll lcm(long long a,long long b){return a/gcd(a,b)*b ;}
// ll sa(long long a,long long b){if(a>b) return a-b ; return b-a ;}
// void fix_cout(){cout << fixed << setprecision(20) ;}

vll fac(2, 1), inv(2, 1), finv(2, 1);

void nCr_set(ll n) {
  for (int i = 2; i <= n; i++) {
    fac.emplace_back(fac.at(fac.size() - 1) * i % mod);
    inv.emplace_back(mod - inv.at(mod % i) * (mod / i) % mod);
    finv.emplace_back(finv.at(finv.size() - 1) * inv.at(i) % mod);
  }
}

ll nCr(ll n, ll r) {
  if (n < 0 || r < 0 || n < r)
    return 0;
  if (n == r || r == 0)
    return 1;
  return fac.at(n) * finv.at(n - r) % mod * finv.at(r) % mod;
}

ll nHr(ll n, ll r) { return nCr(n + r - 1, n - 1); }

void DFS1(vvll &g, vll &sz, int cur, int pre) {
  if (sz.at(cur) != -1)
    return;
  ll sum = 0;
  for (int i = 0; i < g.at(cur).size(); i++) {
    if (pre != g.at(cur).at(i)) {
      DFS1(g, sz, g.at(cur).at(i), cur);
      sum += sz.at(g.at(cur).at(i));
    }
  }
  sz.at(cur) = sum + 1;
}

void DFS2(vvll &g, vll &cnt, vll &sz, int cur) {
  if (cnt.at(cur) != -1)
    return;
  ll sum = 1;
  ll sub = 0;
  for (int i = 0; i < g.at(cur).size(); i++) {
    if (sz.at(cur) > sz.at(g.at(cur).at(i))) {
      DFS2(g, cnt, sz, g.at(cur).at(i));
      sum = sum * nCr(sz.at(cur) - 1 - sub, sz.at(g.at(cur).at(i))) % mod *
            cnt.at(g.at(cur).at(i)) % mod;
      sub += sz.at(g.at(cur).at(i));
    }
  }
  cnt.at(cur) = sum;
}

void DFS3(vvll &g, vll &rev, vll &cnt, vll &sz, int cur, int par, int rem) {
  if (rev.at(cur) != -1)
    return;
  //	cout << cur.first << " " << cur.second << endl ;
  //	cout << cur << endl ;
  ll sum = 1;
  ll sub = 0;
  for (int i = 0; i < g.at(par).size(); i++) {
    if (sz.at(par) < sz.at(g.at(par).at(i))) {
      DFS3(g, rev, cnt, sz, par, g.at(par).at(i), sz.at(0) - sz.at(par));
      sum = sum * nCr(rem - 1 - sub, sz.at(0) - sz.at(par)) % mod *
            rev.at(par) % mod;
      sub += sz.at(0) - sz.at(par);
    } else if (g.at(par).at(i) != cur) {
      sum = sum * nCr(rem - 1 - sub, sz.at(g.at(par).at(i))) % mod *
            cnt.at(g.at(par).at(i)) % mod;
      sub += sz.at(g.at(par).at(i));
    }
  }
  rev.at(cur) = sum;
  //	cout << cur.first << " " << cur.second << " " << rev[cur] << endl ;
}

int main() {
  ll n;
  cin >> n;
  nCr_set(n);
  vvll g(n);
  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    g.at(a).push_back(b);
    g.at(b).push_back(a);
  }
  vll sz(n, -1);
  DFS1(g, sz, 0, -1);
  vll cnt(n, -1), rev(n, -1);
  DFS2(g, cnt, sz, 0);
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < g.at(i).size(); j++) {
      if (sz.at(i) < sz.at(g.at(i).at(j)))
        DFS3(g, rev, cnt, sz, i, g.at(i).at(j), sz.at(0) - sz.at(i));
    }
  }
  for (int i = 0; i < n; i++) {
    ll ans = 1;
    ll sub = 0;
    for (int j = 0; j < g.at(i).size(); j++) {
      if (sz.at(i) > sz.at(g.at(i).at(j))) {
        ans = ans * nCr(n - 1 - sub, sz.at(g.at(i).at(j))) % mod *
              cnt.at(g.at(i).at(j)) % mod;
        sub += sz.at(g.at(i).at(j));
      } else {
        ans = ans * nCr(n - 1 - sub, n - sz.at(i)) % mod * rev.at(i) % mod;
        sub += n - sz.at(i);
      }
    }
    cout << ans << endl;
  }
  //	cout << -1 << endl ;
  //	for(auto i:rev) cout << i << " " ; cout << endl ;
  //	for(auto i:cnt) cout << i << " " ; cout << endl ;
  //	for(auto i:sz) cout << i << " " ; cout << endl ;
}