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

// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
#include <x86intrin.h>
using namespace std;
using ll = long long;
using P = pair<int, int>;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(v) v.begin(), v.end()
#define allr(v) v.rbegin(), v.rend()
#define V vector
template <typename T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename A, size_t N, typename T>
void Fill(A (&array)[N], const T &val) {
  fill((T *)array, (T *)(array + N), val);
}
V<int> dx = {-1, 1, 0, 0, -1, -1, 1, 1};
V<int> dy = {0, 0, -1, 1, -1, 1, -1, 1};

int n;
V<int> a(2e5 + 1, -1), b(2e5 + 1, -1);
V<set<int>> st(2e5 + 1);     // st[i]:i個ある文字のset
map<int, int> ma, mb;        // ma[i]:aに入ってるiの個数
map<int, queue<int>> ia, ib; // ia[i]:aにあるiのindex
V<int> ans(2e5 + 1, -1);

int get_any(map<int, int> &mp, int x) { // x以外をとる
  for (auto p : mp) {
    if (p.first != x)
      return p.first;
  }
  return 0;
}

void use(int x, int y) { // aからx(ia[x][0])を,bからy(ib[y][0])を使う
  int cnt_x = ia[x].size() + ib[x].size();
  st[cnt_x].erase(x);
  if (cnt_x - 1 > 0)
    st[cnt_x - 1].emplace(x);

  int cnt_y = ia[y].size() + ib[y].size();
  st[cnt_y].erase(y);
  if (cnt_y - 1 > 0)
    st[cnt_y - 1].emplace(y);

  ans[ia[x].front()] = y;
  ia[x].pop();
  ib[y].pop();

  ma[x]--;
  mb[y]--;
  if (ma[x] == 0)
    ma.erase(x);
  if (mb[y] == 0)
    mb.erase(y);
}

int main() {
  cin >> n;
  rep(i, n) {
    int a;
    scanf("%d", &a);
    ma[a]++;
    ia[a].emplace(i);
  }
  rep(i, n) {
    int b;
    scanf("%d", &b);
    mb[b]++;
    ib[b].emplace(i);
  }
  for (int i = 1; i <= n; i++) {
    int cnt = ia[i].size() + ib[i].size();
    if (cnt > n) {
      puts("No");
      return 0;
    }
    if (cnt > 0)
      st[cnt].emplace(i);
  }
  puts("Yes");
  for (int i = n; i > 0; i--) {
    if (st[i].size() == 2) { // xとyがi個ずつ
      auto itr = st[i].begin();
      int x = *itr;
      itr++;
      int y = *itr;
      if (ia[x].size() == 0)
        swap(x, y);
      use(x, y);
    } else if (st[i].size() == 1) {
      int z = *st[i].begin(); // 絶対使う
      int x, y;
      if (ia[z].size() > 0) {
        x = z;
        y = get_any(mb, x);
      } else {
        y = z;
        x = get_any(ma, y);
      }
      use(x, y);
    } else {
      int x = get_any(ma, -1);
      int y = get_any(mb, x);
      use(x, y);
    }
  }
  rep(i, n) printf("%d ", ans[i]);
  cout << endl;
  return 0;
}

Error: TLE

Fixed submission:

// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
#include <x86intrin.h>
using namespace std;
using ll = long long;
using P = pair<int, int>;
#define rep(i, n) for (int i = 0; i < (n); i++)
#define all(v) v.begin(), v.end()
#define allr(v) v.rbegin(), v.rend()
#define V vector
template <typename T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename A, size_t N, typename T>
void Fill(A (&array)[N], const T &val) {
  fill((T *)array, (T *)(array + N), val);
}
V<int> dx = {-1, 1, 0, 0, -1, -1, 1, 1};
V<int> dy = {0, 0, -1, 1, -1, 1, -1, 1};

int n;
V<int> a(2e5 + 1, -1), b(2e5 + 1, -1);
V<set<int>> st(2e5 + 1);            // st[i]:i個ある文字のset
map<int, int> ma, mb;               // ma[i]:aに入ってるiの個数
V<V<int>> ia(2e5 + 1), ib(2e5 + 1); // ia[i]:aにあるiのindex
V<int> ans(2e5 + 1, -1);

int get_any(map<int, int> &mp, int x) { // x以外をとる
  for (auto p : mp) {
    if (p.first != x)
      return p.first;
  }
  return 0;
}

void use(int x, int y) { // aからx(ia[x][0])を,bからy(ib[y][0])を使う
  int cnt_x = ia[x].size() + ib[x].size();
  st[cnt_x].erase(x);
  if (cnt_x - 1 > 0)
    st[cnt_x - 1].emplace(x);

  int cnt_y = ia[y].size() + ib[y].size();
  st[cnt_y].erase(y);
  if (cnt_y - 1 > 0)
    st[cnt_y - 1].emplace(y);

  ans[ia[x][0]] = y;
  ia[x].erase(ia[x].begin());
  ib[y].erase(ib[y].begin());

  ma[x]--;
  mb[y]--;
  if (ma[x] == 0)
    ma.erase(x);
  if (mb[y] == 0)
    mb.erase(y);
}

int main() {
  cin >> n;
  rep(i, n) {
    int a;
    scanf("%d", &a);
    ma[a]++;
    ia[a].emplace_back(i);
  }
  rep(i, n) {
    int b;
    scanf("%d", &b);
    mb[b]++;
    ib[b].emplace_back(i);
  }
  for (int i = 1; i <= n; i++) {
    int cnt = ia[i].size() + ib[i].size();
    if (cnt > n) {
      puts("No");
      return 0;
    }
    if (cnt > 0)
      st[cnt].emplace(i);
  }
  puts("Yes");
  for (int i = n; i > 0; i--) {
    if (st[i].size() == 2) { // xとyがi個ずつ
      auto itr = st[i].begin();
      int x = *itr;
      itr++;
      int y = *itr;
      if (ia[x].size() == 0)
        swap(x, y);
      use(x, y);
    } else if (st[i].size() == 1) {
      int z = *st[i].begin(); // 絶対使う
      int x, y;
      if (ia[z].size() > 0) {
        x = z;
        y = get_any(mb, x);
      } else {
        y = z;
        x = get_any(ma, y);
      }
      use(x, y);
    } else {
      int x = get_any(ma, -1);
      int y = get_any(mb, x);
      use(x, y);
    }
  }
  rep(i, n) printf("%d ", ans[i]);
  cout << endl;
  return 0;
}