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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key

typedef long long ll;
typedef pair<ll, ll> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

vector<int> fact;
vector<int> ifact;
vector<int> inv;
vector<int> pow2;
const int MOD = (1e9 + 7);
int add(int a, int b) {
  a += b;
  while (a >= MOD)
    a -= MOD;
  return a;
}
void radd(int &a, int b) { a = add(a, b); }
int mult(int a, int b) { return (a * 1LL * b) % MOD; }
void rmult(int &a, int b) { a = mult(a, b); }
int modpow(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1)
      r = mult(r, a);
    a = mult(a, a);
    b >>= 1;
  }
  return r;
}
int choose(int a, int b) {
  if (a < b)
    return 0;
  if (b == 0)
    return 1;
  if (a == b)
    return 1;
  return mult(fact[a], mult(ifact[b], ifact[a - b]));
}
int inverse(int a) { return modpow(a, MOD - 2); }
void init(int _n) {
  fact.clear();
  ifact.clear();
  inv.clear();
  pow2.clear();
  fact.resize(_n + 1);
  ifact.resize(_n + 1);
  inv.resize(_n + 1);
  pow2.resize(_n + 1);
  pow2[0] = 1;
  ifact[0] = 1;
  fact[0] = 1;
  for (int i = 1; i <= _n; i++) {
    pow2[i] = add(pow2[i - 1], pow2[i - 1]);
    fact[i] = mult(fact[i - 1], i);
    // ifact[i]=mult(ifact[i-1],inv[i]);
  }
  ifact[_n] = inverse(fact[_n]);
  for (int i = _n - 1; i >= 1; i--) {
    ifact[i] = mult(ifact[i + 1], i + 1);
  }
  for (int i = 1; i <= _n; i++) {
    inv[i] = mult(fact[i - 1], ifact[i]);
  }
}

int a[5555];
bool vis[5555];
int poww[5555];
int w[5555];
vi adj[5555];

void dfs(int u, int r) {
  w[r]++;
  vis[u] = 1;
  for (int v : adj[u]) {
    if (!vis[v])
      dfs(v, r);
  }
}

int dp[5555][5555];

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  init(88888);
  int n;
  cin >> n;
  int emp = 0;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    if (a[i] > 0) {
      a[i]--;
      adj[i].pb(a[i]);
      adj[a[i]].pb(i);
    } else
      emp++;
  }
  poww[0] = 1;
  for (int i = 1; i <= 5002; i++)
    poww[i] = mult(poww[i - 1], n - 1);
  int cyc = 0;
  // count # of cycles
  for (int i = 0; i < n; i++) {
    if (!vis[i]) {
      int cur = i;
      int ex = 0;
      set<int> S;
      while (a[cur] != -1 && !vis[a[cur]]) {
        cur = a[cur];
        if (S.find(cur) != S.end())
          break;
        S.insert(cur);
        if (cur == i) {
          ex = 1;
          break;
        }
      }
      if (ex) {
        cyc++;
        cur = i;
        vis[cur] = 1;
        while (a[cur] != i) {
          cur = a[cur];
          vis[cur] = 1;
        }
      }
    }
  }
  memset(vis, 0, sizeof(vis));
  for (int i = 0; i < n; i++) {
    if (a[i] == -1) {
      w[i] = 0;
      dfs(i, i);
      // cerr<<i<<' '<<w[i]<<'\\n';
    }
  }
  cyc = mult(cyc, poww[emp]);
  /*
  for(int i=1;i<=n;i++) //length of cycle
  {
          cyc=add(cyc,mult(poww[emp-i],fact[i-1]));
  }
  */
  vi V;
  for (int i = 0; i < n; i++) {
    if (a[i] == -1)
      V.pb(w[i]);
  }
  for (int v : V) {
    // cerr<<v<<' ';
    cyc = add(cyc, mult(poww[emp - 1], v - 1));
  }
  // cerr<<cyc<<'\\n';
  // cerr<<'\\n';
  dp[0][0] = 1;
  for (int i = 1; i <= V.size(); i++) {
    for (int j = 1; j <= V.size(); j++) {
      if (j == 1)
        dp[i][j] = add(dp[i - 1][j], V[i - 1]);
      else
        dp[i][j] = add(dp[i - 1][j], mult(dp[i - 1][j - 1], V[i - 1]));
    }
  }
  for (int i = 2; i <= V.size(); i++) {
    // cerr<<i<<' '<<dp[int(V.size())][i]<<'\\n';
    cyc =
        add(cyc, mult(poww[emp - i], mult(fact[i - 1], dp[int(V.size())][i])));
  }
  int ans = add(mult(n, poww[emp]), MOD - cyc);
  cout << ans << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key

typedef long long ll;
typedef pair<ll, ll> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

vector<int> fact;
vector<int> ifact;
vector<int> inv;
vector<int> pow2;
const int MOD = (1e9 + 7);
int add(int a, int b) {
  a += b;
  while (a >= MOD)
    a -= MOD;
  return a;
}
void radd(int &a, int b) { a = add(a, b); }
int mult(int a, int b) { return (a * 1LL * b) % MOD; }
void rmult(int &a, int b) { a = mult(a, b); }
int modpow(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1)
      r = mult(r, a);
    a = mult(a, a);
    b >>= 1;
  }
  return r;
}
int choose(int a, int b) {
  if (a < b)
    return 0;
  if (b == 0)
    return 1;
  if (a == b)
    return 1;
  return mult(fact[a], mult(ifact[b], ifact[a - b]));
}
int inverse(int a) { return modpow(a, MOD - 2); }
void init(int _n) {
  fact.clear();
  ifact.clear();
  inv.clear();
  pow2.clear();
  fact.resize(_n + 1);
  ifact.resize(_n + 1);
  inv.resize(_n + 1);
  pow2.resize(_n + 1);
  pow2[0] = 1;
  ifact[0] = 1;
  fact[0] = 1;
  for (int i = 1; i <= _n; i++) {
    pow2[i] = add(pow2[i - 1], pow2[i - 1]);
    fact[i] = mult(fact[i - 1], i);
    // ifact[i]=mult(ifact[i-1],inv[i]);
  }
  ifact[_n] = inverse(fact[_n]);
  for (int i = _n - 1; i >= 1; i--) {
    ifact[i] = mult(ifact[i + 1], i + 1);
  }
  for (int i = 1; i <= _n; i++) {
    inv[i] = mult(fact[i - 1], ifact[i]);
  }
}

int a[5555];
bool vis[5555];
int poww[5555];
int w[5555];
vi adj[5555];

void dfs(int u, int r) {
  w[r]++;
  vis[u] = 1;
  for (int v : adj[u]) {
    if (!vis[v])
      dfs(v, r);
  }
}

int dp[5555][5555];

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  init(88888);
  int n;
  cin >> n;
  int emp = 0;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    if (a[i] > 0) {
      a[i]--;
      adj[i].pb(a[i]);
      adj[a[i]].pb(i);
    } else
      emp++;
  }
  poww[0] = 1;
  for (int i = 1; i <= 5002; i++)
    poww[i] = mult(poww[i - 1], n - 1);
  int cyc = 0;
  // count # of cycles
  for (int i = 0; i < n; i++) {
    if (!vis[i]) {
      int cur = i;
      int ex = 0;
      while (a[cur] != -1 && !vis[a[cur]]) {
        cur = a[cur];
        if (cur == i) {
          ex = 1;
          break;
        }
      }
      if (ex) {
        cyc++;
        cur = i;
        vis[cur] = 1;
        while (a[cur] != i) {
          cur = a[cur];
          vis[cur] = 1;
        }
      }
    }
  }
  memset(vis, 0, sizeof(vis));
  for (int i = 0; i < n; i++) {
    if (a[i] == -1) {
      w[i] = 0;
      dfs(i, i);
      // cerr<<i<<' '<<w[i]<<'\\n';
    }
  }
  cyc = mult(cyc, poww[emp]);
  /*
  for(int i=1;i<=n;i++) //length of cycle
  {
          cyc=add(cyc,mult(poww[emp-i],fact[i-1]));
  }
  */
  vi V;
  for (int i = 0; i < n; i++) {
    if (a[i] == -1)
      V.pb(w[i]);
  }
  for (int v : V) {
    // cerr<<v<<' ';
    cyc = add(cyc, mult(poww[emp - 1], v - 1));
  }
  // cerr<<cyc<<'\\n';
  // cerr<<'\\n';
  dp[0][0] = 1;
  for (int i = 1; i <= V.size(); i++) {
    for (int j = 1; j <= V.size(); j++) {
      if (j == 1)
        dp[i][j] = add(dp[i - 1][j], V[i - 1]);
      else
        dp[i][j] = add(dp[i - 1][j], mult(dp[i - 1][j - 1], V[i - 1]));
    }
  }
  for (int i = 2; i <= V.size(); i++) {
    // cerr<<i<<' '<<dp[int(V.size())][i]<<'\\n';
    cyc =
        add(cyc, mult(poww[emp - i], mult(fact[i - 1], dp[int(V.size())][i])));
  }
  int ans = add(mult(n, poww[emp]), MOD - cyc);
  cout << ans << '\\n';
}