AtCoder が提供しているABC（AtCoder Beginner Contest）329 のF問題をC++とPythonで解いてみました。ABC329は、2023年11月18日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
F問題 Colored Ball（Difficulty : 1207）
問題はリンク先をご覧ください。
setコンテナを工夫して使います。AtCoder Problems による Difficulty は 1207 でした。
解答案
最初に TLE（Time Limit Exceeded）判定となるプログラムを紹介します。以下は、補足とプログラムです。
- 箱に入るボールの色を格納する set コンテナ c を用意します。
- それぞれのクエリ (a, b) に対して以下を行う。
- c[a] にある要素を c[b] に加える。c[a] は clear する（20ー24行目）。
- c[b] の大きさを出力する（25行目）。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, q;
cin >> n >> q;
vector<set<int>> c(n);
for (int i = 0; i < n; ++i) {
int t;
cin >> t;
c[i].insert(t);
}
for (int i = 0; i < q; ++i) {
int a, b;
cin >> a >> b;
--a;
--b;
for (auto e : c[a]) {
c[b].insert(e);
}
c[a].clear();
cout << c[b].size() << endl;
}
return 0;
}このプログラムは、テストケース46件中、6件で TLE 判定でした（2023年11月19日の環境）。ここで以下の工夫をします。
箱aと箱bの色が少ない方から多い方に移す。箱bからaに移動した場合は、箱を入れ替える。
これにより、移した結果は少ない箱の2倍以上になるため、移動する計算量は $O(\log n)$ となり間に合います。
C++ プログラム例（ABC329F）
上記の TLE したプログラムから以下を変更します。
- c[a] より c[b] の方が小さい場合は、同じ処理となります（20ー24行目）。
- 上記でない場合は、c[b] から c[a] に要素（色）を移して、最後に swap する（26ー30行目）。
差分は小さいですが、このプログラムは AC（Accepted＝正しいプログラム）と判定されました。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, q;
cin >> n >> q;
vector<set<int>> c(n);
for (int i = 0; i < n; ++i) {
int t;
cin >> t;
c[i].insert(t);
}
for (int i = 0; i < q; ++i) {
int a, b;
cin >> a >> b;
--a;
--b;
if (c[a].size() < c[b].size()) {
for (auto e : c[a]) {
c[b].insert(e);
}
c[a].clear();
} else {
for (auto e : c[b]) {
c[a].insert(e);
}
c[b].clear();
swap(c[a], c[b]);
}
cout << c[b].size() << endl;
}
return 0;
}Python プログラム例（ABC329F）
Python の set は順序がありませんが、この問題は順序が関係ないため、同じように書くことができます。以下となります。
"""AtCoder Beginner Contest 329 F"""
n, q = map(int, input().split())
t = list(map(int, input().split()))
c = [set([t[i]]) for i in range(n)]
for i in range(q):
a, b = map(int, input().split())
a -= 1
b -= 1
if len(c[a]) < len(c[b]):
for e in c[a]:
c[b].add(e)
c[a].clear()
else:
for e in c[b]:
c[a].add(e)
c[b].clear()
c[b], c[a] = c[a], c[b]
print(len(c[b]))こちらも「AC」と判定されました。
最後に
TLE判定のプログラムとAC判定のプログラムの差は、「要素数の移動が少ない方から多い方に移動する」だけです。「マージテク」と呼ばれているようです。
コンテスト中に、この工夫に気付ければ、気分が良いと思います。
引き続き ABC の問題を紹介していきます。
