Problem p02809 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
#define vv(a, b, c, d) vector<vector<d>>(a, vector<d>(b, c))
#define vvi vector<vector<int>>
#define vvl vector<vector<ll>>
#define vll vector<ll>
#define rep(c, a, b) for (ll c = a; c < b; c++)
#define re(c, b) for (ll c = 0; c < b; c++)
#define all(obj) (obj).begin(), (obj).end()
typedef long long int ll;
typedef long double ld;
using namespace std;

//-------------------------------------------------------------------
struct segtree_in {
  ll M = 1;
  vector<vll> dat;
  segtree_in(ll N, vll num) {
    while (M < N)
      M *= 2;
    for (int i = 0; i < M * 2 - 1; i++)
      dat.push_back(num);
  }
  void update(vll x, int k, int l = 0, int r = -1) {
    if (r == -1)
      r = M;
    k += M - 1;
    dat[k] = x;
    while (k > 0) {
      k = (k - 1) / 2;
      dat[k] = (dat[k * 2 + 1][0] > dat[k * 2 + 2][0] ? dat[k * 2 + 1]
                                                      : dat[k * 2 + 2]);
    }
  }
  vll query(int a, int b = -1, int k = 0, int l = 0, int r = -1) {
    if (r == -1)
      r = M;
    if (b == -1)
      b = M;
    if (r <= a || b <= l)
      return vll{-1000000000, -1};
    if (a <= l && r <= b)
      return dat[k];
    vll A = query(a, b, k * 2 + 1, l, (l + r) / 2);
    vll B = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return (A[0] > B[0] ? A : B);
  }
};
//-------------------------------------------------------------------
int main(int argc, char const *argv[]) {
  ll n;
  std::cin >> n;
  vll a(n + 1), uke(n + 1, 0), ans;
  re(i, n) {
    std::cin >> a[i + 1];
    uke[a[i + 1]]++;
  }
  segtree_in segin(n + 1, vll{-1000000000, -1});
  set<ll> s;
  rep(i, 1, n + 1) {
    segin.update(vll{uke[i], i}, i);
    s.insert(i);
  }

  if (n == 2) {
    std::cout << -1 << '\n';
    return 0;
  }
  for (int i = 0; i < n; i++) {
    vll k = segin.query(1, n + 1);
    ll p;
    if (k[0] == (n - 1 - i) && (i == 0 || a[ans[i - 1]] != k[1]))
      p = k[1];
    else {
      auto itr = s.begin();
      if (i != 0 && a[ans[i - 1]] == *itr)
        itr++;
      if (itr == s.end())
        while (true)
          n++;
      p = *itr;
    }
    segin.update(vll{-1000000000, -1}, p);
    if (i > 0) {
      ll to = a[ans[i - 1]];
      vll g = segin.query(to, to + 1);
      segin.update(vll{g[0] - 1, to}, to);
    }

    ans.push_back(p);
    s.erase(p);
  }
  re(i, n) std::cout << ans[i] << (i == n - 1 ? "\n" : " ");
  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
#define vv(a, b, c, d) vector<vector<d>>(a, vector<d>(b, c))
#define vvi vector<vector<int>>
#define vvl vector<vector<ll>>
#define vll vector<ll>
#define rep(c, a, b) for (ll c = a; c < b; c++)
#define re(c, b) for (ll c = 0; c < b; c++)
#define all(obj) (obj).begin(), (obj).end()
typedef long long int ll;
typedef long double ld;
using namespace std;

//-------------------------------------------------------------------
struct segtree_in {
  ll M = 1;
  vector<vll> dat;
  segtree_in(ll N, vll num) {
    while (M < N)
      M *= 2;
    for (int i = 0; i < M * 2 - 1; i++)
      dat.push_back(num);
  }
  void update(vll x, int k, int l = 0, int r = -1) {
    if (r == -1)
      r = M;
    k += M - 1;
    dat[k] = x;
    while (k > 0) {
      k = (k - 1) / 2;
      dat[k] = (dat[k * 2 + 1][0] > dat[k * 2 + 2][0] ? dat[k * 2 + 1]
                                                      : dat[k * 2 + 2]);
    }
  }
  vll query(int a, int b = -1, int k = 0, int l = 0, int r = -1) {
    if (r == -1)
      r = M;
    if (b == -1)
      b = M;
    if (r <= a || b <= l)
      return vll{-1000000000, -1};
    if (a <= l && r <= b)
      return dat[k];
    vll A = query(a, b, k * 2 + 1, l, (l + r) / 2);
    vll B = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return (A[0] > B[0] ? A : B);
  }
};
//-------------------------------------------------------------------
int main(int argc, char const *argv[]) {
  ll n;
  std::cin >> n;
  vll a(n + 1), uke(n + 1, 0), ans;
  re(i, n) {
    std::cin >> a[i + 1];
    uke[a[i + 1]]++;
  }
  segtree_in segin(n + 1, vll{-1000000000, -1});
  set<ll> s;
  rep(i, 1, n + 1) {
    segin.update(vll{uke[i], i}, i);
    s.insert(i);
  }

  if (n == 2) {
    std::cout << -1 << '\n';
    return 0;
  }
  for (int i = 0; i < n; i++) {
    vll k = segin.query(1, n + 1);
    ll p;
    if (k[0] == (n - 1 - i) && (i == 0 || a[ans[i - 1]] != k[1]))
      p = k[1];
    else {
      if (s.size() == 3) {
        for (auto t = s.begin(); t != s.end(); t++) {
          p = *t;
          ll A = 10000000000, B = 0;
          for (auto x = s.begin(); x != s.end(); x++)
            if (x != t)
              A = min(A, *x), B = max(B, *x);
          if ((a[A] == B && a[B] == A) || (i != 0 && a[ans[i - 1]] == p))
            continue;
          break;
        }
      } else {
        auto itr = s.begin();
        if (i != 0 && a[ans[i - 1]] == *itr)
          itr++;
        p = *itr;
      }
    }
    segin.update(vll{-1000000000, -1}, p);
    if (i > 0) {
      ll to = a[ans[i - 1]];
      vll g = segin.query(to, to + 1);
      segin.update(vll{g[0] - 1, to}, to);
    }

    ans.push_back(p);
    s.erase(p);
  }
  re(i, n) std::cout << ans[i] << (i == n - 1 ? "\n" : " ");
  return 0;
}