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

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

struct DS {
  vector<int> t;
  DS(int n) {
    t.resize(n + 1);
    for (int i = 1; i <= n; i++) {
      t[i] = i;
    }
  }
  int root(int i) {
    if (t[i] == i) {
      return i;
    } else {
      return t[i] = root(t[i]);
    }
  }
  void unite(int i, int j) { t[root(i)] = root(j); }
  int cnt() {
    int sol = 0;
    for (int i = 1; i < (int)t.size(); i++) {
      if (t[i] == i) {
        sol++;
      }
    }
    return sol;
  }
};

const int N = 5000 + 7;
const int M = (int)1e9 + 7;

int add(int a, int b) {
  a += b;
  if (a >= M) {
    return a - M;
  }
  if (a < 0) {
    return a + M;
  }
  return a;
}

int mul(int a, int b) { return a * (ll)b % M; }

int pw(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1) {
      r = mul(r, a);
    }
    a = mul(a, a);
    b /= 2;
  }
  return r;
}

int n;
int p[N];

void read() {
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
  }
}

int nodesDfs;
int edgesDfs;
vector<int> g[N];
bool vis[N];

void dfs(int a) {
  vis[a] = 1;
  nodesDfs++;
  if (p[a] != -1) {
    edgesDfs++;
  }
  for (auto &b : g[a]) {
    if (vis[b] == 0) {
      dfs(b);
    }
  }
}

vector<int> b;
int universal2;
int sum[N][N], el[N];
int yolo[N];

void bkt2(int pos, vector<int> orz) {
  if (pos == (int)b.size()) {
    if ((int)orz.size() <= 1)
      return;
    int remunsolved = 0;
    remunsolved = (int)b.size() - (int)orz.size();
    int rem = pw(n - 1, remunsolved);
    int now;
    now = 1;
    for (auto &it : orz) {
      now = mul(now, it);
    }
    for (int i = 1; i < (int)orz.size(); i++) {
      now = mul(now, i);
    }
    ///    yolo[(int) orz.size()] = add(yolo[(int) orz.size(), ]);
    universal2 = add(universal2, mul(rem, now));
    return;
  }
  bkt2(pos + 1, orz);
  orz.push_back(b[pos]);
  bkt2(pos + 1, orz);
}

int smart() {
  int cntTotal = 1;
  for (int i = 1; i <= n; i++) {
    if (p[i] == -1) {
      cntTotal = mul(cntTotal, n - 1);
    }
  }
  for (int i = 1; i <= n; i++) {
    if (p[i] != -1) {
      g[i].push_back(p[i]);
      g[p[i]].push_back(i);
    }
  }
  int sol = 0;
  for (int i = 1; i <= n; i++) {
    if (vis[i] == 0) {
      nodesDfs = 0;
      edgesDfs = 0;
      dfs(i);
      if (edgesDfs == nodesDfs) {
        sol = add(sol, cntTotal);
      }
      if (edgesDfs == nodesDfs - 1) {
        b.push_back(nodesDfs);
      }
    }
  }
  for (auto &it : b) {
    sol = add(sol, mul(pw(n - 1, (int)b.size() - 1), it - 1));
  }
  vector<int> fact(N);
  fact[0] = 1;
  for (int i = 1; i < N; i++) {
    fact[i] = mul(fact[i - 1], i);
  }
  int cnt = (int)b.size();
  for (int i = 1; i <= cnt; i++) {
    el[i] = b[i - 1];
  }
  for (int i = 1; i <= cnt; i++) {
    sum[i][1] = add(sum[i - 1][1], el[i]);
    for (int j = 2; j <= i; j++) {
      sum[i][j] = sum[i - 1][j];
      int val = mul(sum[i - 1][j - 1], el[i]);
      sum[i][j] = add(sum[i][j], val);
    }
  }
  int altern = 0;
  for (int c = 2; c <= cnt; c++) {
    int mlt = fact[c - 1];
    int rem = pw(n - 1, cnt - c);
    mlt = mul(mlt, rem);
    altern = add(altern, mul(sum[cnt][c], mlt));
  }
  return add(sol, altern);
}

int solve() {
  int sol = n;
  for (int i = 1; i <= n; i++) {
    if (p[i] == -1) {
      sol = mul(sol, n - 1);
    }
  }
  sol = add(sol, -smart());
  return sol;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  read();
  cout << solve() << "\\n";
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

struct DS {
  vector<int> t;
  DS(int n) {
    t.resize(n + 1);
    for (int i = 1; i <= n; i++) {
      t[i] = i;
    }
  }
  int root(int i) {
    if (t[i] == i) {
      return i;
    } else {
      return t[i] = root(t[i]);
    }
  }
  void unite(int i, int j) { t[root(i)] = root(j); }
  int cnt() {
    int sol = 0;
    for (int i = 1; i < (int)t.size(); i++) {
      if (t[i] == i) {
        sol++;
      }
    }
    return sol;
  }
};

const int N = 5000 + 7;
const int M = (int)1e9 + 7;

int add(int a, int b) {
  a += b;
  if (a >= M) {
    return a - M;
  }
  if (a < 0) {
    return a + M;
  }
  return a;
}

int mul(int a, int b) { return a * (ll)b % M; }

int pw(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1) {
      r = mul(r, a);
    }
    a = mul(a, a);
    b /= 2;
  }
  return r;
}

int n;
int p[N];

void read() {
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
  }
}

int nodesDfs;
int edgesDfs;
vector<int> g[N];
bool vis[N];

void dfs(int a) {
  vis[a] = 1;
  nodesDfs++;
  if (p[a] != -1) {
    edgesDfs++;
  }
  for (auto &b : g[a]) {
    if (vis[b] == 0) {
      dfs(b);
    }
  }
}

vector<int> b;
int universal2;

void bkt2(int pos, vector<int> orz) {
  if (pos == (int)b.size()) {
    if ((int)orz.size() <= 1)
      return;
    int remunsolved = 0;
    remunsolved = (int)b.size() - (int)orz.size();
    int rem = pw(n - 1, remunsolved);
    int now;
    now = 1;
    for (auto &it : orz) {
      now = mul(now, it);
    }
    for (int i = 1; i < (int)orz.size(); i++) {
      now = mul(now, i);
    }
    universal2 = add(universal2, mul(rem, now));
    return;
  }
  bkt2(pos + 1, orz);
  orz.push_back(b[pos]);
  bkt2(pos + 1, orz);
}

int smart() {
  int cntTotal = 1;
  for (int i = 1; i <= n; i++) {
    if (p[i] == -1) {
      cntTotal = mul(cntTotal, n - 1);
    }
  }
  for (int i = 1; i <= n; i++) {
    if (p[i] != -1) {
      g[i].push_back(p[i]);
      g[p[i]].push_back(i);
    }
  }
  int sol = 0;
  for (int i = 1; i <= n; i++) {
    if (vis[i] == 0) {
      nodesDfs = 0;
      edgesDfs = 0;
      dfs(i);
      if (edgesDfs == nodesDfs) {
        sol = add(sol, cntTotal);
      }
      if (edgesDfs == nodesDfs - 1) {
        b.push_back(nodesDfs);
      }
    }
  }
  for (auto &it : b) {
    sol = add(sol, mul(pw(n - 1, (int)b.size() - 1), it - 1));
  }
  universal2 = 0;
  bkt2(0, {});
  return add(sol, universal2);
}

int solve() {
  int sol = n;
  for (int i = 1; i <= n; i++) {
    if (p[i] == -1) {
      sol = mul(sol, n - 1);
    }
  }
  sol = add(sol, -smart());
  return sol;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  read();
  cout << solve() << "\\n";
  return 0;
}