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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
#define all(x) (x).begin(), (x).end()
#define fi first
#define se second
#define mp make_pair
#define si(x) int(x.size())
const int mod = 1000000007, MAX = 100005, INF = 1 << 30;

int main() {

  std::ifstream in("text.txt");
  std::cin.rdbuf(in.rdbuf());
  cin.tie(0);
  ios::sync_with_stdio(false);

  int N;
  cin >> N;
  set<pair<int, int>> sum;
  vector<int> A(N), B(N), ans(N), ca(N + 1), cb(N + 1), cc(N + 1);
  for (int i = 0; i < N; i++) {
    cin >> A[i];
    ca[A[i]]++;
    cc[A[i]]++;
  }
  for (int i = 0; i < N; i++) {
    cin >> B[i];
    cb[B[i]]++;
    cc[B[i]]++;
  }

  for (int i = 1; i <= N; i++) {
    if (cc[i] && cb[i])
      sum.insert(mp(cc[i], i));

    if (cc[i] > N) {
      cout << "No" << endl;
      return 0;
    }
  }

  cout << "Yes" << endl;

  for (int i = 0; i < N; i++) {
    auto it = sum.end();
    it--;

    while ((*it).se == A[i])
      it--;

    ans[i] = (*it).se;

    sum.erase(*it);
    cc[ans[i]]--;
    cb[ans[i]]--;
    if (cc[ans[i]] && cb[ans[i]])
      sum.insert(mp(cc[ans[i]], ans[i]));

    sum.erase(mp(cc[A[i]], A[i]));
    cc[A[i]]--;
    ca[A[i]]--;
    if (cc[A[i]])
      sum.insert(mp(cc[A[i]], A[i]));
  }

  for (int i = 0; i < N; i++)
    cout << ans[i] << " ";
  cout << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
#define all(x) (x).begin(), (x).end()
#define fi first
#define se second
#define mp make_pair
#define si(x) int(x.size())
const int mod = 1000000007, MAX = 100005, INF = 1 << 30;

int main() {

  std::ifstream in("text.txt");
  std::cin.rdbuf(in.rdbuf());
  cin.tie(0);
  ios::sync_with_stdio(false);

  int N;
  cin >> N;
  set<pair<int, int>> cnt1, cnt2, sum;
  vector<int> A(N), B(N), ans(N), ca(N + 1), cb(N + 1), cc(N + 1);
  for (int i = 0; i < N; i++) {
    cin >> A[i];
    ca[A[i]]++;
    cc[A[i]]++;
  }
  for (int i = 0; i < N; i++) {
    cin >> B[i];
    cb[B[i]]++;
    cc[B[i]]++;
  }

  for (int i = 1; i <= N; i++) {
    if (ca[i])
      cnt1.insert(mp(ca[i], i));
    if (cb[i])
      cnt2.insert(mp(cb[i], i));
    if (cc[i] && cb[i])
      sum.insert(mp(cc[i], i));

    if (cc[i] > N) {
      cout << "No" << endl;
      return 0;
    }
  }

  cout << "Yes" << endl;

  for (int i = 0; i < N; i++) {
    auto it = sum.end();
    it--;

    while ((*it).se == A[i] || cb[(*it).se] == 0)
      it--;

    ans[i] = (*it).se;

    sum.erase(*it);
    cc[ans[i]]--;
    sum.insert(mp(cc[ans[i]], ans[i]));

    cb[ans[i]]--;

    sum.erase(mp(cc[A[i]], A[i]));
    cc[A[i]]--;
    sum.insert(mp(cc[A[i]], A[i]));

    ca[A[i]]--;
  }

  for (int i = 0; i < N; i++)
    cout << ans[i] << " ";
  cout << endl;
}