Problem p03650 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, b) for (int i = (a); i <= (b); ++i)
#define REPD(i, a, b) for (int i = (a); i >= (b); --i)
#define FORI(i, n) REP(i, 1, n)
#define FOR(i, n) REP(i, 0, int(n) - 1)
#define mp make_pair
#define pb push_back
#define pii pair<int, int>
#define vi vector<int>
#define ll long long
#define SZ(x) int((x).size())
#define DBG(v) cerr << #v << " = " << (v) << endl;
#define FOREACH(i, t) for (typeof(t.begin()) i = t.begin(); i != t.end(); i++)
#define fi first
#define se second

const int N = 200200;

int n;
int vis[N], val[N], outdeg[N], p[N];
set<int> vs[N];
int la;

int mex(set<int> &s) {
  set<int>::iterator it = s.begin();
  int ret = 0;
  while (it != s.end()) {
    if (*it != ret)
      break;
    it++;
    ret++;
  }
  return ret;
}

void dfs(int u) {
  vis[u] = 1;
  val[u] = mex(vs[u]);
  vs[p[u]].insert(val[u]);
  outdeg[p[u]]--;
  if (outdeg[p[u]] == 0)
    dfs(p[u]);
}

bool ok(int u, int vv) {
  int v = u;
  val[u] = vv;
  while (p[v] != u) {
    set<int> st = vs[p[v]];
    st.insert(val[v]);
    int vc = mex(st);
    val[p[v]] = vc;
    v = p[v];
  }
  la = v;
  return val[u] != val[v];
}

int main() {
  scanf("%d", &n);
  FOR(i, n) scanf("%d", &p[i]);
  FOR(i, n) p[i]--;
  FOR(i, n) outdeg[p[i]]++;
  FOR(i, n) if (!vis[i] && outdeg[i] == 0) dfs(i);
  FOR(i, n) if (!vis[i]) {
    int v1 = mex(vs[i]);
    if (!ok(i, v1)) {
      vs[i].insert(v1);
      int v2 = mex(vs[i]);
      if (!ok(i, v2) || val[la] != v1) {
        printf("IMPOSSIBLE\n");
        return 0;
      }
    }
  }
  printf("POSSIBLE\n");
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, b) for (int i = (a); i <= (b); ++i)
#define REPD(i, a, b) for (int i = (a); i >= (b); --i)
#define FORI(i, n) REP(i, 1, n)
#define FOR(i, n) REP(i, 0, int(n) - 1)
#define mp make_pair
#define pb push_back
#define pii pair<int, int>
#define vi vector<int>
#define ll long long
#define SZ(x) int((x).size())
#define DBG(v) cerr << #v << " = " << (v) << endl;
#define FOREACH(i, t) for (typeof(t.begin()) i = t.begin(); i != t.end(); i++)
#define fi first
#define se second

const int N = 200200;

int n;
int vis[N], val[N], outdeg[N], p[N];
set<int> vs[N];
int la;

int mex(set<int> &s) {
  set<int>::iterator it = s.begin();
  int ret = 0;
  while (it != s.end()) {
    if (*it != ret)
      break;
    it++;
    ret++;
  }
  return ret;
}

void dfs(int u) {
  vis[u] = 1;
  val[u] = mex(vs[u]);
  vs[p[u]].insert(val[u]);
  outdeg[p[u]]--;
  if (outdeg[p[u]] == 0)
    dfs(p[u]);
}

bool ok(int u, int vv) {
  int v = u;
  val[u] = vv;
  while (p[v] != u) {
    set<int> st = vs[p[v]];
    st.insert(val[v]);
    int vc = mex(st);
    val[p[v]] = vc;
    vis[v] = 1;
    v = p[v];
  }
  la = v;
  return val[u] != val[v];
}

int main() {
  scanf("%d", &n);
  FOR(i, n) scanf("%d", &p[i]);
  FOR(i, n) p[i]--;
  FOR(i, n) outdeg[p[i]]++;
  FOR(i, n) if (!vis[i] && outdeg[i] == 0) dfs(i);
  FOR(i, n) if (!vis[i]) {
    int v1 = mex(vs[i]);
    if (!ok(i, v1)) {
      vs[i].insert(v1);
      int v2 = mex(vs[i]);
      if (!ok(i, v2) || val[la] != v1) {
        printf("IMPOSSIBLE\n");
        return 0;
      }
    }
  }
  printf("POSSIBLE\n");
  return 0;
}