Problem p02557 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#define fi first
#define se second
#define mp make_pair
using namespace std;

typedef pair<int, int> pii;

set<pii> s1, s2;
int a[200010], b[200010], ans[200010];
int pos[200010], cnt[200010];

set<pii>::iterator get_pre(set<pii>::iterator it) {
  it--;
  return it;
}

int main() {
  int n;
  scanf("%d", &n);
  memset(pos, -1, sizeof(pos));
  for (int i = 0; i < n; i++) {
    scanf("%d", &a[i]);
    if (pos[a[i]] == -1)
      pos[a[i]] = i;
    cnt[a[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i])
      s1.insert(mp(cnt[i], i));
  }
  for (int i = 0; i < n; i++) {
    scanf("%d", &b[i]);
    cnt[b[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i] > n) {
      printf("No\n");
      return 0;
    }
  }
  memset(cnt, 0, sizeof(cnt));
  for (int i = 0; i < n; i++) {
    cnt[b[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i])
      s2.insert(mp(cnt[i], i));
  }
  for (int i = 0; i < n; i++) {
    int cnt = 0;
    for (auto x : s1)
      cnt += x.fi;
    set<pii>::iterator it = s1.end();
    it--;
    if (it == s1.begin() || get_pre(it)->fi != it->fi) {
      pii tmp = *it;
      s1.erase(tmp);
      tmp.fi--;
      if (tmp.fi)
        s1.insert(tmp);
      it = s2.lower_bound(mp(get_pre(s2.end())->fi, -1));
      if (it->se == tmp.se) {
        it++;
        if (it == s2.end())
          it--, it--;
      }
      pii ttt = *it;
      s2.erase(ttt);
      ttt.fi--;
      if (ttt.fi)
        s2.insert(ttt);
      ans[pos[tmp.se]++] = ttt.se;
    } else {
      it = s2.end();
      it--;
      pii tmp = *it;
      s2.erase(tmp);
      tmp.fi--;
      if (tmp.fi)
        s2.insert(tmp);
      it = s1.lower_bound(mp(get_pre(s1.end())->fi, -1));
      if (it->se == tmp.se) {
        it++;
        if (it == s1.end()) {
          it--, it--;
        }
      }
      pii ttt = *it;
      s1.erase(ttt);
      ttt.fi--;
      if (ttt.fi)
        s1.insert(ttt);
      ans[pos[ttt.se]++] = tmp.se;
    }
  }
  printf("Yes\n");
  for (int i = 0; i < n; i++) {
    printf("%d ", ans[i]);
  }
  printf("\n");
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define fi first
#define se second
#define mp make_pair
using namespace std;

typedef pair<int, int> pii;

set<pii> s1, s2;
int a[200010], b[200010], ans[200010];
int pos[200010], cnt[200010];

set<pii>::iterator get_pre(set<pii>::iterator it) {
  it--;
  return it;
}

int main() {
  int n;
  scanf("%d", &n);
  memset(pos, -1, sizeof(pos));
  for (int i = 0; i < n; i++) {
    scanf("%d", &a[i]);
    if (pos[a[i]] == -1)
      pos[a[i]] = i;
    cnt[a[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i])
      s1.insert(mp(cnt[i], i));
  }
  for (int i = 0; i < n; i++) {
    scanf("%d", &b[i]);
    cnt[b[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i] > n) {
      printf("No\n");
      return 0;
    }
  }
  memset(cnt, 0, sizeof(cnt));
  for (int i = 0; i < n; i++) {
    cnt[b[i]]++;
  }
  for (int i = 1; i <= n; i++) {
    if (cnt[i])
      s2.insert(mp(cnt[i], i));
  }
  for (int i = 0; i < n; i++) {
    set<pii>::iterator it = s1.end();
    it--;
    if (it == s1.begin() || get_pre(it)->fi != it->fi) {
      pii tmp = *it;
      s1.erase(tmp);
      tmp.fi--;
      if (tmp.fi)
        s1.insert(tmp);
      it = s2.lower_bound(mp(get_pre(s2.end())->fi, -1));
      if (it->se == tmp.se) {
        it++;
        if (it == s2.end())
          it--, it--;
      }
      pii ttt = *it;
      s2.erase(ttt);
      ttt.fi--;
      if (ttt.fi)
        s2.insert(ttt);
      ans[pos[tmp.se]++] = ttt.se;
    } else {
      it = s2.end();
      it--;
      pii tmp = *it;
      s2.erase(tmp);
      tmp.fi--;
      if (tmp.fi)
        s2.insert(tmp);
      it = s1.lower_bound(mp(get_pre(s1.end())->fi, -1));
      if (it->se == tmp.se) {
        it++;
        if (it == s1.end()) {
          it--, it--;
        }
      }
      pii ttt = *it;
      s1.erase(ttt);
      ttt.fi--;
      if (ttt.fi)
        s1.insert(ttt);
      ans[pos[ttt.se]++] = tmp.se;
    }
  }
  printf("Yes\n");
  for (int i = 0; i < n; i++) {
    printf("%d ", ans[i]);
  }
  printf("\n");
  return 0;
}