Problem p03650 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <map>
#include <memory>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

#define ws ws_____________________
#define y1 y1_____________________
#define y0 y0_____________________
#define left left_________________
#define right right_______________
#define next next_________________
#define prev prev_________________
#define hash hash_________________

#define pb push_back
#define fst first
#define snd second
#define mp make_pair
#define sz(C) ((int)(C).size())
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define ford(i, n) for (int i = int(n) - 1; i >= 0; --i)
#define all(C) begin(C), end(C)

typedef long long ll;
typedef unsigned long long ull;
typedef unsigned int uint;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<pii> vii;
typedef long double ld;
typedef complex<double> cd;

#ifdef LOCAL
#define eprintf(args...) fprintf(stderr, args), fflush(stderr)
#else
#define eprintf(...) ;
#endif

#define FILE_NAME "a"

int n;
vvi g;

bool read() {
  if (scanf("%d", &n) < 1) {
    return 0;
  }
  g.clear();
  g.resize(n);
  forn(i, n) {
    int p;
    scanf("%d", &p);
    --p;
    g[p].pb(i);
  }
  return 1;
}

bool dfs_cycle(int v, vi &st, vi &col, vi &ans) {
  col[v] = 1;
  st.pb(v);

  for (int to : g[v]) {
    if (col[to] == 2) {
      continue;
    }

    if (col[to] == 0) {
      if (dfs_cycle(to, st, col, ans)) {
        return 1;
      }
    }

    assert(col[to] == 1);
    ford(i, sz(st)) {
      assert(col[st[i]] == 1);
      ans.pb(st[i]);
      if (st[i] == to) {
        break;
      }
    }

    reverse(all(ans));
    return 1;
  }

  col[v] = 2;
  assert(st.back() == v);
  st.pop_back();
  return 0;
}

void dfs(int v, vi &val) {
  set<int> vals;
  for (int to : g[v]) {
    dfs(to, val);
    vals.insert(val[to]);
  }

  for (int x = 0;; ++x) {
    if (!vals.count(x)) {
      val[v] = x;
      break;
    }
  }
}

bool check(const vi &cycle, vi &val) {
  ford(i, sz(cycle)) {
    const int v = cycle[i];
    if (val[v] != -1) {
      continue;
    }

    set<int> vals;
    for (int to : g[v]) {
      assert(val[to] != -1);
      vals.insert(val[to]);
    }

    for (int x = 0;; ++x) {
      if (!vals.count(x)) {
        val[v] = x;
        break;
      }
    }
  }

  for (int v : cycle) {
    set<int> vals;
    for (int to : g[v]) {
      vals.insert(val[to]);
    }

    if (vals.count(val[v])) {
      return 0;
    }

    for (int x = 0;; ++x) {
      if (!vals.count(x)) {
        if (val[v] != x) {
          return 0;
        }
        break;
      }
    }
  }

  return 1;
}

bool solve() {
  forn(v, n) {
    for (int to : g[v]) {
      if (to == v) {
        return 0;
      }
    }
  }

  vi col(n, 0);
  vi st;
  vi cycle;
  forn(v, n) {
    if (!col[v]) {
      if (dfs_cycle(v, st, col, cycle)) {
        break;
      }
    }
  }
  st.clear();
  col.assign(n, 0);

  assert(!cycle.empty());

  // eprintf("cycle:\n");
  // for (int v : cycle) {
  // 	eprintf("%d ", v + 1);
  // }
  // eprintf("\n");

  vi on_cycle(n, 0);
  for (int v : cycle) {
    on_cycle[v] = 1;
  }

  vi val(n, -1);
  forn(v, n) {
    if (on_cycle[v]) {
      for (int to : g[v]) {
        if (!on_cycle[to]) {
          dfs(to, val);
        }
      }
    }
  }

  int root = cycle[0];
  set<int> sons;
  for (int to : g[root]) {
    if (!on_cycle[to]) {
      sons.insert(val[to]);
    }
  }

  int cnt = 0;
  for (int x = 0;; ++x) {
    if (!sons.count(x)) {
      val[root] = x;

      if (check(cycle, val)) {
        return 1;
      }

      for (int u : cycle) {
        val[u] = -1;
      }

      ++cnt;

      if (cnt == 2) {
        return 0;
      }
    }
  }

  return 0;
}

int main() {
#ifdef LOCAL
  freopen(FILE_NAME ".in", "r", stdin);
  // freopen(FILE_NAME ".out", "w", stdout);
#endif

  while (read()) {
    puts(solve() ? "POSSIBLE" : "IMPOSSIBLE");
  }

#ifdef LOCAL
  cerr.precision(5);
  cerr << "Time: " << fixed << (double)clock() / CLOCKS_PER_SEC << endl;
#endif
  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <map>
#include <memory>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

#define ws ws_____________________
#define y1 y1_____________________
#define y0 y0_____________________
#define left left_________________
#define right right_______________
#define next next_________________
#define prev prev_________________
#define hash hash_________________

#define pb push_back
#define fst first
#define snd second
#define mp make_pair
#define sz(C) ((int)(C).size())
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define ford(i, n) for (int i = int(n) - 1; i >= 0; --i)
#define all(C) begin(C), end(C)

typedef long long ll;
typedef unsigned long long ull;
typedef unsigned int uint;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<ll> vll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<pii> vii;
typedef long double ld;
typedef complex<double> cd;

#ifdef LOCAL
#define eprintf(args...) fprintf(stderr, args), fflush(stderr)
#else
#define eprintf(...) ;
#endif

#define FILE_NAME "a"

int n;
vvi g;

bool read() {
  if (scanf("%d", &n) < 1) {
    return 0;
  }
  g.clear();
  g.resize(n);
  forn(i, n) {
    int p;
    scanf("%d", &p);
    --p;
    g[p].pb(i);
  }
  return 1;
}

bool dfs_cycle(int v, vi &st, vi &col, vi &ans) {
  col[v] = 1;
  st.pb(v);

  for (int to : g[v]) {
    if (col[to] == 2) {
      continue;
    }

    if (col[to] == 0) {
      if (dfs_cycle(to, st, col, ans)) {
        return 1;
      }
      continue;
    }

    assert(col[to] == 1);
    ford(i, sz(st)) {
      assert(col[st[i]] == 1);
      ans.pb(st[i]);
      if (st[i] == to) {
        break;
      }
    }

    reverse(all(ans));
    return 1;
  }

  col[v] = 2;
  assert(st.back() == v);
  st.pop_back();
  return 0;
}

void dfs(int v, vi &val) {
  set<int> vals;
  for (int to : g[v]) {
    dfs(to, val);
    vals.insert(val[to]);
  }

  for (int x = 0;; ++x) {
    if (!vals.count(x)) {
      val[v] = x;
      break;
    }
  }
}

bool check(const vi &cycle, vi &val) {
  ford(i, sz(cycle)) {
    const int v = cycle[i];
    if (val[v] != -1) {
      continue;
    }

    set<int> vals;
    for (int to : g[v]) {
      assert(val[to] != -1);
      vals.insert(val[to]);
    }

    for (int x = 0;; ++x) {
      if (!vals.count(x)) {
        val[v] = x;
        break;
      }
    }
  }

  for (int v : cycle) {
    set<int> vals;
    for (int to : g[v]) {
      vals.insert(val[to]);
    }

    if (vals.count(val[v])) {
      return 0;
    }

    for (int x = 0;; ++x) {
      if (!vals.count(x)) {
        if (val[v] != x) {
          return 0;
        }
        break;
      }
    }
  }

  return 1;
}

bool solve() {
  forn(v, n) {
    for (int to : g[v]) {
      if (to == v) {
        return 0;
      }
    }
  }

  vi col(n, 0);
  vi st;
  vi cycle;
  forn(v, n) {
    if (!col[v]) {
      if (dfs_cycle(v, st, col, cycle)) {
        break;
      }
    }
  }
  st.clear();
  col.assign(n, 0);

  assert(!cycle.empty());

  // eprintf("cycle:\n");
  // for (int v : cycle) {
  // 	eprintf("%d ", v + 1);
  // }
  // eprintf("\n");

  vi on_cycle(n, 0);
  for (int v : cycle) {
    on_cycle[v] = 1;
  }

  vi val(n, -1);
  forn(v, n) {
    if (on_cycle[v]) {
      for (int to : g[v]) {
        if (!on_cycle[to]) {
          dfs(to, val);
        }
      }
    }
  }

  int root = cycle[0];
  set<int> sons;
  for (int to : g[root]) {
    if (!on_cycle[to]) {
      sons.insert(val[to]);
    }
  }

  int cnt = 0;
  for (int x = 0;; ++x) {
    if (!sons.count(x)) {
      val[root] = x;

      if (check(cycle, val)) {
        return 1;
      }

      for (int u : cycle) {
        val[u] = -1;
      }

      ++cnt;

      if (cnt == 2) {
        return 0;
      }
    }
  }

  return 0;
}

int main() {
#ifdef LOCAL
  freopen(FILE_NAME ".in", "r", stdin);
  // freopen(FILE_NAME ".out", "w", stdout);
#endif

  while (read()) {
    puts(solve() ? "POSSIBLE" : "IMPOSSIBLE");
  }

#ifdef LOCAL
  cerr.precision(5);
  cerr << "Time: " << fixed << (double)clock() / CLOCKS_PER_SEC << endl;
#endif
  return 0;
}