ABC 170 E - Smart Infants

ABC 170 E - Smart Infants についてのメモを残します。

multiset

• multiset - cpprefjp C++日本語リファレンス
• AOJ ITP2_7_D Multi-Set

set は要素の重複を許さないが、multiset は重複を許すデータ構造となっている。

#include <bits/stdc++.h>
using namespace std;

int main() {
  multiset<int> S = { 1, 2, 2, 3, 3, 3 };
  cout << S.size() << endl;  // 6
  S.erase(3);
  cout << S.size() << endl;  // 3
  S.erase(S.find(2));
  cout << S.size() << endl;  // 2
}

erase 関数に引数として要素を渡すと、その要素が全て削除されている。イテレータを渡すと、イテレータが指す要素 1 個のみが削除されている。

PAST？

• 問題 : ABC 170 E - Smart Infants

下図は、入力例 1 における初期状態。

• 各幼稚園の最大レートを高速に取得する
• 最強園児たちの最低レート (平等さ) を高速に取得する

を実現したい。これらは set を用いればよさそう。
2 種類 set を用意する。1 つは、各幼稚園のレートの集合 (200,000 個の set)。もう 1 つは、最強園児たちのレートの集合。
ただし、1 つの幼稚園に同レートの幼児が複数人 所属する状況や、異なる幼稚園の最強園児のレートが同一となる状況があり得るので、両方とも要素の重複を許す multiset で実装する。

さて、クエリ処理は「幼児 を幼稚園 へ転園させる」といった内容。

• 幼児 がどの幼稚園に所属しているか
• 幼児 のレートはいくらか

という情報を保持しておく必要がある。
下図は、入力例 1 における「幼児 4 を幼稚園 3 へ転園させる」の様子を表したもの。転園前、幼児 4 は幼稚園 1 に所属しており、レートは 1 である。

この図で「変わるかも」と指摘している箇所があるが、この転園では値は変わっていない。
実装でやることは以下の通り。

• 幼稚園 1 と幼稚園 3 の最大のレート を取得し、最強園児たちのレートの集合から を削除。ここでは
• 幼稚園 1 のレートの集合から 1 を削除し、幼稚園 3 のレートの集合に 1 を追加
• 幼稚園 1 と幼稚園 3 の最大のレート を取得し、最強園児たちのレートの集合に を追加。ここでは
• 幼児 4 が所属する幼稚園の番号を 3 に更新
• 最強園児たちのレートの集合の最小値 (平等さ) を出力

誰も所属していない幼稚園の存在や、multiset の erase 関数の挙動に気をつけながら実装していく。

#include <bits/stdc++.h>
using namespace std;

int main() {
  int N, Q; cin >> N >> Q;
  vector<int> A(N), B(N);
  vector<multiset<int>> R(200000);  // 幼稚園 i のレートの集合
  multiset<int> M;  // 最強園児たちのレートの集合
  for (int i = 0; i < N; i++) {
    cin >> A.at(i) >> B.at(i);
    B.at(i)--;
    R.at(B.at(i)).insert(A.at(i));
  }
  for (int i = 0; i < 200000; i++) {
    if (!R.at(i).empty()) M.insert(*R.at(i).rbegin());
  }

  for (int i = 0; i < Q; i++) {
    int c, d; cin >> c >> d;
    c--; d--;
    int prev = B.at(c);  // 転園前

    // M から幼稚園 prev と幼稚園 d の最大レートを削除
    M.erase(M.find(*R.at(prev).rbegin()));
    if (!R.at(d).empty()) M.erase(M.find(*R.at(d).rbegin()));

    // 幼稚園 prev のレートの集合から A[c] を削除
    R.at(prev).erase(R.at(prev).find(A.at(c)));

    // 幼稚園 d のレートの集合に A[c] を追加
    R.at(d).insert(A.at(c));

    // M に幼稚園 prev と幼稚園 d の最大レートを追加
    if (!R.at(prev).empty()) M.insert(*R.at(prev).rbegin());
    M.insert(*R.at(d).rbegin());

    B.at(c) = d;  // 所属情報の更新

    cout << *M.begin() << endl;
  }
}