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

#include <bits/stdc++.h>
#include <stdlib.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vec;
typedef vector<vec> mat;
typedef pair<ll, ll> P;
typedef priority_queue<P, vector<P>, greater<P>> P_queue;

const ll MOD = 998244353;
const ll mod = 1000000007;
const ll INF = 1e18;
const double PI = 3.1415926535897932;
vec dx = {1, 0, -1, 0};
vec dy = {0, 1, 0, -1};

#define REP(i, a, b) for (int i = (int)a; i < (int)b; i++)
#define rep(i, n) REP(i, 0, n)
#define pb push_back
#define mp make_pair
#define ALL(a) a.begin(), a.end()
#define SORT(a) sort(ALL(a))
#define U_ERASE(V) V.erase(unique(ALL(V)), V.end());
#define ADD(a, b) a = (a + b) % mod
#define MUL(a, b) a = (a * b) % mod

const ll MAXN = 200011;
vector<P> d[MAXN];
ll D[2 * MAXN][2];
ll N;

ll memo[2 * MAXN], siz[2 * MAXN];
vec Lm[MAXN], Rm[MAXN], Ls[MAXN], Rs[MAXN];

const ll MaxN = 201001;

ll f[MaxN], rf[MaxN];
ll inv(ll x) {
  ll res = 1;
  ll k = mod - 2;
  ll y = x;
  while (k) {
    if (k & 1)
      res = (res * y) % mod;
    y = (y * y) % mod;
    k /= 2;
  }
  return res;
}
void init() {
  f[0] = 1;
  REP(i, 1, MaxN) f[i] = (f[i - 1] * i) % mod;
  REP(i, 0, MaxN) rf[i] = inv(f[i]);
}
//---------------------------------------------------------------------------------------------------
ll C(ll n, ll k) {
  ll a = f[n];      // = n!
  ll b = rf[n - k]; // = (n-k)!
  ll c = rf[k];     // = k!

  ll bc = (b * c) % mod;

  return (a * bc) % mod;
}

ll Ai(ll x) {
  if (x % 2)
    return x - 1;
  return x + 1;
}

ll Size(ll x) {
  if (siz[x] > 0)
    return siz[x];
  ll A = D[x][0], B = D[x][1];
  if (d[B].size() == 1)
    return siz[x] = 1;
  ll ret = 1;
  rep(i, d[B].size()) if (d[B][i].second != Ai(x)) ret += Size(d[B][i].second);
  return siz[x] = ret;
}
ll Solve(ll x) {
  if (memo[x] > 0)
    return memo[x];
  ll A = D[x][0], B = D[x][1];
  if (d[B].size() == 1)
    return memo[x] = 1;
  P hikaku = mp(A, -INF);
  ll K = upper_bound(ALL(d[B]), hikaku) - d[B].begin();
  ll LM = Lm[B].size(), RM = Rm[B].size();
  // cout<<x<<' '<<K<<endl;
  // if(x==529) return 0;
  if (LM <= K) {
    REP(j, LM - 1, K) {
      ll kari = Lm[B][j], p = d[B][j].second;
      Ls[B].pb(Size(p) + Ls[B][j]);
      MUL(kari, Solve(p));
      MUL(kari, C(Ls[B][j + 1], Ls[B][j]));
      Lm[B].pb(kari);
    }
  }
  if (RM < d[B].size() - K) {
    REP(j, RM - 1, d[B].size() - K - 1) {
      ll kari = Rm[B][j], p = d[B][d[B].size() - j - 1].second;
      Rs[B].pb(Size(p) + Rs[B][j]);
      MUL(kari, Solve(p));
      MUL(kari, C(Rs[B][j + 1], Rs[B][j]));
      Rm[B].pb(kari);
    }
  }
  ll ret = 1;

  MUL(ret, Lm[B][K]);
  MUL(ret, Rm[B][d[B].size() - K - 1]);
  MUL(ret, C(Rs[B][d[B].size() - K - 1] + Ls[B][K], Ls[B][K]));
  return memo[x] = ret;
}

ll SOLVE(ll B) {
  if (d[B].size() == 1)
    return Solve(d[B][0].second);
  ll ret = 1;
  ll X = d[B][0].second, Y = Ai(X);
  MUL(ret, C(N - 1, Size(X)));
  MUL(ret, Solve(X));
  MUL(ret, Solve(Y));
  return ret;
}

int main() {
  init();
  rep(i, MAXN) {
    Lm[i].pb(1);
    Rm[i].pb(1);
    Ls[i].pb(0);
    Rs[i].pb(0);
  }
  cin >> N;
  rep(i, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    d[a].pb(mp(b, 2 * i));
    d[b].pb(mp(a, 2 * i + 1));
    D[2 * i][0] = a;
    D[2 * i][1] = b;
    D[2 * i + 1][0] = b;
    D[2 * i + 1][1] = a;
  }
  rep(i, N) sort(ALL(d[i]));
  // rep(i,2*N) cout<<Size(i)<<endl;
  rep(i, N) cout << SOLVE(i) << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <stdlib.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vec;
typedef vector<vec> mat;
typedef pair<ll, ll> P;
typedef priority_queue<P, vector<P>, greater<P>> P_queue;

const ll MOD = 998244353;
const ll mod = 1000000007;
const ll INF = 1e18;
const double PI = 3.1415926535897932;
vec dx = {1, 0, -1, 0};
vec dy = {0, 1, 0, -1};

#define REP(i, a, b) for (int i = (int)a; i < (int)b; i++)
#define rep(i, n) REP(i, 0, n)
#define pb push_back
#define mp make_pair
#define ALL(a) a.begin(), a.end()
#define SORT(a) sort(ALL(a))
#define U_ERASE(V) V.erase(unique(ALL(V)), V.end());
#define ADD(a, b) a = (a + b) % mod
#define MUL(a, b) a = (a * b) % mod

const ll MAXN = 200011;
vector<P> d[MAXN];
ll D[2 * MAXN][2];
ll N;

ll memo[2 * MAXN], siz[2 * MAXN];
vec Lm[MAXN], Rm[MAXN], Ls[MAXN], Rs[MAXN];

const ll MaxN = 201001;

ll f[MaxN], rf[MaxN];
ll inv(ll x) {
  ll res = 1;
  ll k = mod - 2;
  ll y = x;
  while (k) {
    if (k & 1)
      res = (res * y) % mod;
    y = (y * y) % mod;
    k /= 2;
  }
  return res;
}
void init() {
  f[0] = 1;
  REP(i, 1, MaxN) f[i] = (f[i - 1] * i) % mod;
  REP(i, 0, MaxN) rf[i] = inv(f[i]);
}
//---------------------------------------------------------------------------------------------------
ll C(ll n, ll k) {
  ll a = f[n];      // = n!
  ll b = rf[n - k]; // = (n-k)!
  ll c = rf[k];     // = k!

  ll bc = (b * c) % mod;

  return (a * bc) % mod;
}

ll Ai(ll x) {
  if (x % 2)
    return x - 1;
  return x + 1;
}

ll Size(ll x) {
  if (siz[x] > 0)
    return siz[x];
  ll A = D[x][0], B = D[x][1];
  if (d[B].size() == 1)
    return siz[x] = 1;
  ll ret = 1;
  rep(i, d[B].size()) if (d[B][i].second != Ai(x)) ret += Size(d[B][i].second);
  return siz[x] = ret;
}
ll Solve(ll x) {
  if (memo[x] > 0)
    return memo[x];
  ll A = D[x][0], B = D[x][1];
  if (d[B].size() == 1)
    return memo[x] = 1;
  P hikaku = mp(A, -INF);
  ll K = upper_bound(ALL(d[B]), hikaku) - d[B].begin();
  ll LM = Lm[B].size(), RM = Rm[B].size();
  // cout<<x<<' '<<K<<endl;
  // if(x==529) return 0;
  if (LM <= K) {
    REP(j, LM - 1, K) {
      ll kari = Lm[B][j], p = d[B][j].second;
      Ls[B].pb(Size(p) + Ls[B][j]);
      MUL(kari, Solve(p));
      MUL(kari, C(Ls[B][j + 1], Ls[B][j]));
      Lm[B].pb(kari);
    }
  }
  if (RM < d[B].size() - K) {
    REP(j, RM - 1, d[B].size() - K - 1) {
      ll kari = Rm[B][j], p = d[B][d[B].size() - j - 1].second;
      Rs[B].pb(Size(p) + Rs[B][j]);
      MUL(kari, Solve(p));
      MUL(kari, C(Rs[B][j + 1], Rs[B][j]));
      Rm[B].pb(kari);
    }
  }
  ll ret = 1;

  MUL(ret, Lm[B][K]);
  MUL(ret, Rm[B][d[B].size() - K - 1]);
  MUL(ret, C(Rs[B][d[B].size() - K - 1] + Ls[B][K], Ls[B][K]));
  return memo[x] = ret;
}

ll SOLVE(ll a) {
  ll M = N - 1;
  ll ret = 1;
  rep(i, d[a].size()) {
    ll c = d[a][i].second;
    MUL(ret, Solve(c));
    MUL(ret, C(M, Size(c)));
    M -= (Size(c));
  }
  return ret;
}

int main() {
  init();
  rep(i, MAXN) {
    Lm[i].pb(1);
    Rm[i].pb(1);
    Ls[i].pb(0);
    Rs[i].pb(0);
  }
  cin >> N;
  rep(i, N - 1) {
    ll a, b;
    cin >> a >> b;
    a--;
    b--;
    d[a].pb(mp(b, 2 * i));
    d[b].pb(mp(a, 2 * i + 1));
    D[2 * i][0] = a;
    D[2 * i][1] = b;
    D[2 * i + 1][0] = b;
    D[2 * i + 1][1] = a;
  }
  rep(i, N) sort(ALL(d[i]));
  // rep(i,2*N) cout<<Size(i)<<endl;
  rep(i, N) cout << SOLVE(i) << endl;
}