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

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int Maxn = 5005;
const int mod = 1000000007;

int pw[Maxn];
int n;
int a[Maxn];
int par[Maxn], siz[Maxn];
bool spec[Maxn];
vector<int> seq;
int dp[Maxn];
int res;

int getPar(int x) { return par[x] == x ? x : par[x] = getPar(par[x]); }

void unionSet(int a, int b) {
  a = getPar(a), b = getPar(b);
  if (a == b)
    return;
  if (siz[a] < siz[b])
    swap(a, b);
  siz[a] += siz[b];
  par[b] = a;
  spec[a] |= spec[b];
}

int main() {
  scanf("%d", &n);
  pw[0] = 1;
  for (int i = 1; i < Maxn; i++)
    pw[i] = ll(pw[i - 1]) * (n - 1) % mod;
  for (int i = 1; i <= n; i++) {
    par[i] = i;
    siz[i] = 1;
    scanf("%d", &a[i]);
    spec[i] = a[i] == -1;
  }
  for (int i = 1; i <= n; i++)
    if (a[i] != -1)
      unionSet(i, a[i]);
  for (int i = 1; i <= n; i++)
    if (a[i] == -1)
      seq.push_back(siz[getPar(i)]);
  for (int i = 1; i <= n; i++)
    if (getPar(i) == i && !spec[i])
      res = (res + pw[seq.size()]) % mod;
  sort(seq.begin(), seq.end());
  for (int i = 0; i < seq.size(); i++) {
    for (int j = i; j >= 1; j--)
      dp[j + 1] = (dp[j + 1] + ll(dp[j]) * j % mod * seq[i]) % mod;
    dp[1] = (dp[1] + seq[i]) % mod;
  }
  for (int i = 0; i < seq.size(); i++)
    res = (res + ll(seq[i] - 1) * pw[int(seq.size()) - 1]) % mod;
  for (int j = 2; j <= seq.size(); j++)
    res = (res + ll(dp[j]) * pw[int(seq.size()) - j]) % mod;
  res = (mod - res) % mod;
  res = (res + ll(n) * pw[seq.size()]) % mod;
  printf("%d\\n", res);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int Maxn = 5005;
const int mod = 1000000007;

int pw[Maxn];
int n;
int a[Maxn];
int par[Maxn], siz[Maxn];
bool spec[Maxn];
vector<int> seq;
int dp[Maxn][Maxn];
int res;

int getPar(int x) { return par[x] == x ? x : par[x] = getPar(par[x]); }

void unionSet(int a, int b) {
  a = getPar(a), b = getPar(b);
  if (a == b)
    return;
  if (siz[a] < siz[b])
    swap(a, b);
  siz[a] += siz[b];
  par[b] = a;
  spec[a] |= spec[b];
}

int main() {
  scanf("%d", &n);
  pw[0] = 1;
  for (int i = 1; i < Maxn; i++)
    pw[i] = ll(pw[i - 1]) * (n - 1) % mod;
  for (int i = 1; i <= n; i++) {
    par[i] = i;
    siz[i] = 1;
    scanf("%d", &a[i]);
    spec[i] = a[i] == -1;
  }
  for (int i = 1; i <= n; i++)
    if (a[i] != -1)
      unionSet(i, a[i]);
  for (int i = 1; i <= n; i++)
    if (a[i] == -1)
      seq.push_back(siz[getPar(i)]);
  for (int i = 1; i <= n; i++)
    if (getPar(i) == i && !spec[i])
      res = (res + pw[seq.size()]) % mod;
  sort(seq.begin(), seq.end());
  int sum = 0;
  for (int i = 0; i < seq.size(); i++) {
    for (int j = i; j >= 1; j--)
      for (int k = 0; k <= sum; k++)
        if (dp[j][k])
          dp[j + 1][k + seq[i]] =
              (dp[j + 1][k + seq[i]] + ll(dp[j][k]) * j % mod * seq[i]) % mod;
    dp[1][seq[i]] = (dp[1][seq[i]] + seq[i]) % mod;
    sum += seq[i];
  }
  for (int i = 0; i < seq.size(); i++)
    res = (res + ll(seq[i] - 1) * pw[int(seq.size()) - 1]) % mod;
  for (int j = 2; j <= seq.size(); j++)
    for (int k = 1; k <= sum; k++)
      if (dp[j][k])
        res = (res + ll(dp[j][k]) * pw[int(seq.size()) - j]) % mod;
  res = (mod - res) % mod;
  res = (res + ll(n) * pw[seq.size()]) % mod;
  printf("%d\\n", res);
  return 0;
}