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

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define MOD2 998244353
#define PRIM 3
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int fact[1100000];
Int revfact[1100000];
Int rev[1100000];

Int mod_pow(Int x, Int a, Int m) {
  if (a == 0)
    return 1;
  Int res = mod_pow(x, a / 2, m);
  res = res * res % m;
  if (a % 2)
    res *= x;
  return res % m;
}

Int inv(Int x, Int m) { return mod_pow(x, m - 2, m); }

void init() {
  Int m = MOD;
  fact[0] = 1;
  revfact[0] = 1;
  rev[0] = 0;
  rev[1] = 1;
  for (int i = 1; i < 1100000; i++) {
    fact[i] = fact[i - 1] * i % m;
    if (i > 1)
      rev[i] = MOD / i * (MOD - rev[MOD % i]) % MOD;
    revfact[i] = revfact[i - 1] * rev[i] % MOD;
  }
}

Int nCk(Int n, Int k) {
  if (n < k)
    return 0;
  return fact[n] * revfact[n - k] % MOD * revfact[k] % MOD;
}

class UnionFind {
public:
  int n;
  vector<int> u;
  UnionFind(int _n = 0) {
    n = _n;
    u = vector<int>(n, -1);
  }
  int r(int x) {
    if (u[x] < 0)
      return x;
    return u[x] = r(u[x]);
  }
  void unite(int x, int y) {
    x = r(x);
    y = r(y);
    if (x == y)
      return;
    if (-u[x] < -u[y])
      swap(x, y);
    u[x] += u[y];
    u[y] = x;
  }
  int size(int x) { return -u[r(x)]; }
};

Int pwn[5500];
int p[5500];
Int dp[5500];
vector<int> root;
int cnt;
Int ans;
int main() {
  init();
  int n;
  cin >> n;
  UnionFind uf(n);
  for (int i = 0; i < n; i++) {
    cin >> p[i];
    if (p[i] != -1) {
      p[i]--;
      if (uf.r(i) != uf.r(p[i])) {
        cnt++;
        uf.unite(i, p[i]);
      }
    } else {
      root.push_back(i);
    }
  }
  ans = (root.size() + cnt) * mod_pow(n - 1, root.size(), MOD) % MOD;
  dp[0] = 1;
  pwn[0] = 1;
  for (int i = 1; i < 5500; i++)
    pwn[i] = pwn[i - 1] * (n - 1) % MOD;
  for (auto x : root) {
    for (int i = 0; i < root.size(); i++) {
      Int t = uf.size(x);
      if (i == 0)
        t--;
      ans -= dp[i] * fact[i] % MOD * t % MOD * pwn[root.size() - i - 1] % MOD;
      ans %= MOD;
    }
    for (int i = n; i > 0; i--) {
      dp[i] += dp[i - 1] * uf.size(x) % MOD;
      dp[i] %= MOD;
    }
  }
  if (ans < 0)
    ans += MOD;
  cout << ans << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define MOD2 998244353
#define PRIM 3
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int fact[1100000];
Int revfact[1100000];
Int rev[1100000];

Int mod_pow(Int x, Int a, Int m) {
  if (a == 0)
    return 1;
  Int res = mod_pow(x, a / 2, m);
  res = res * res % m;
  if (a % 2)
    res *= x;
  return res % m;
}

Int inv(Int x, Int m) { return mod_pow(x, m - 2, m); }

void init() {
  Int m = MOD;
  fact[0] = 1;
  revfact[0] = 1;
  rev[0] = 0;
  rev[1] = 1;
  for (int i = 1; i < 1100000; i++) {
    fact[i] = fact[i - 1] * i % m;
    if (i > 1)
      rev[i] = MOD / i * (MOD - rev[MOD % i]) % MOD;
    revfact[i] = revfact[i - 1] * rev[i] % MOD;
  }
}

Int nCk(Int n, Int k) {
  if (n < k)
    return 0;
  return fact[n] * revfact[n - k] % MOD * revfact[k] % MOD;
}

class UnionFind {
public:
  int n;
  vector<int> u;
  UnionFind(int _n = 0) {
    n = _n;
    u = vector<int>(n, -1);
  }
  int r(int x) {
    if (u[x] < 0)
      return x;
    return u[x] = r(u[x]);
  }
  void unite(int x, int y) {
    x = r(x);
    y = r(y);
    if (x == y)
      return;
    if (-u[x] < -u[y])
      swap(x, y);
    u[x] += u[y];
    u[y] = x;
  }
  int size(int x) { return -u[r(x)]; }
};

int p[5500];
Int dp[5500];
vector<int> root;
int cnt;
Int ans;
int main() {
  init();
  int n;
  cin >> n;
  UnionFind uf(n);
  for (int i = 0; i < n; i++) {
    cin >> p[i];
    if (p[i] != -1) {
      p[i]--;
      if (uf.r(i) != uf.r(p[i])) {
        cnt++;
        uf.unite(i, p[i]);
      }
    } else {
      root.push_back(i);
    }
  }
  ans = (root.size() + cnt) * mod_pow(n - 1, root.size(), MOD) % MOD;
  dp[0] = 1;
  for (auto x : root) {
    for (int i = 0; i < root.size(); i++) {
      Int t = uf.size(x);
      if (i == 0)
        t--;
      ans -= dp[i] * fact[i] % MOD * t % MOD *
             mod_pow(n - 1, root.size() - i - 1, MOD) % MOD;
      ans %= MOD;
    }
    for (int i = n; i > 0; i--) {
      dp[i] += dp[i - 1] * uf.size(x) % MOD;
      dp[i] %= MOD;
    }
  }
  if (ans < 0)
    ans += MOD;
  cout << ans << endl;
  return 0;
}