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

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>
#define pq priority_queue
using namespace std;

int main() {
  //	ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#ifdef LOCAL
  //		freopen("in1.txt", "r", stdin);
  //		freopen("in1.txt", "w", stdout);
#endif
  int n;
  cin >> n;
  vector<int> a(n), b(n);
  vector<int> ca(n + 1), cb(n + 1);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    ca[a[i]]++;
  }
  for (int i = 0; i < n; i++) {
    cin >> b[i];
    cb[b[i]]++;
  }
  bool poss = true;
  for (int i = 1; i <= n && poss; i++) {
    if (ca[i] + cb[i] > n)
      poss = false;
  }
  if (!poss) {
    cout << "No\\n"; // LOWERCASE
    return 0;
  }
  cout << "Yes\\n";

  vector<int> at(n + 1);
  for (int i = 0; i < n; i++) {
    if (i && a[i] == a[i - 1])
      continue;
    at[a[i]] = i;
  }
  pq<pii> pa, pb;
  for (int i = 1; i <= n; i++) {
    if (ca[i])
      pa.push({ca[i], i});
    if (cb[i])
      pb.push({cb[i], i});
  }
  vector<int> ans(n);
  while (!pa.empty() && !pb.empty()) {
    int fa = pa.top().first;
    int ia = pa.top().second;
    pa.pop();
    int fb = pb.top().first;
    int ib = pb.top().second;
    pb.pop();
    if (ia == ib) {
      if (!pb.empty()) {
        int tfb = pb.top().first;
        int tib = pb.top().second;
        pb.pop();
        pb.push({fb, ib});
        fb = tfb;
        ib = tib;
      } else if (!pa.empty()) {
        int tfa = pa.top().first;
        int tia = pa.top().second;
        pa.pop();
        pa.push({fa, ia});
        fa = tfa;
        ia = tia;
      }
    }
    //		cout << fa << ' ' << ia << ' ' << fb << ' ' << ib << ' ' << at[ia] <<
    //'\\n';
    int take = min(fa, fb);
    fa -= take;
    fb -= take;
    while (take--) {
      ans[at[ia]] = ib;
      at[ia]++;
    }
    if (fa)
      pa.push({fa, ia});
    if (fb)
      pb.push({fb, ib});
  }

  int idx = -1;
  for (int i = 0; i < n; i++) {
    if (ans[i] == a[i]) {
      idx = i;
      break;
    }
  }
  if (idx >= 0) {
    for (int j = 0; j < n && idx < n && ans[idx] == a[idx]; j++) {
      if (ans[j] != a[idx] && ans[idx] != a[j]) {
        swap(ans[j], ans[idx]);
        idx++;
      }
    }
  }
  for (int x : ans)
    cout << x << ' ';
  cout << '\\n';
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>
using namespace std;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
#ifdef LOCAL
  //		freopen("in1.txt", "r", stdin);
  //		freopen("in1.txt", "w", stdout);
#endif
  int n;
  cin >> n;
  vector<int> a(n), b(n);
  vector<int> ca(n + 1), cb(n + 1);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    ca[a[i]]++;
  }
  for (int i = 0; i < n; i++) {
    cin >> b[i];
    cb[b[i]]++;
  }
  bool poss = true;
  for (int i = 1; i <= n && poss; i++) {
    if (ca[i] + cb[i] > n)
      poss = false;
  }
  if (!poss) {
    cout << "No\\n"; // LOWERCASE
    return 0;
  }
  cout << "Yes\\n";

  vector<bool> vis(n);
  vector<int> ans(n);
  int ida = 0, idb = 0;
  while (ida < n) {
    while (vis[idb] || a[ida] == b[idb])
      idb = (idb + 1) % n;
    vis[idb] = true;
    ans[ida] = b[idb];
    ida++;
    idb = (idb + 1) % n;
  }
  for (int x : ans)
    cout << x << ' ';
  cout << '\\n';
}