ABC330 E問題（Mex and Update）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）330 のE問題をC++とPythonで解いてみました。ABC330は、2023年11月25日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

E問題 Mex and Update（Difficulty : 1004）

問題はリンク先をご覧ください。

ABC330 E問題 Mex and Update

setコンテナを使ってMEXの値を更新します。AtCoder Problems による Difficulty は 1004 でした。

解答案

N個の要素からなる集合のmex（その集合に含まれない最小の非負整数）について、以下が成立します。

N個の要素からなる集合のmexは、N以下になる。

これは、0からNまでの (N+1) 個の要素のうち、最大N個までしか存在しないためです。このため、N以下だけ考えれば（更新していけば）良いことが分かります。

C++ プログラム例（ABC330E）

setコンテナmexは、以下とします。s[i]は0からnまでの数字の頻度を格納します。

• 0からnまでの数字で、a[i]としてでてこない要素をすべてmexに格納しておく（18ー22行目）。
• a[i]の値をxに更新します。
  • s[a[i]]の値をひとつ減らして、その数字がなくなればa[i]の値をmexに加えます（29ー32行目）。
  • xがいままで存在していなければ、xの値をmexから消去して、s[x]の値を更新します（36ー39行目）。
  • 毎回のクエリの結果として、mexに格納されている最小値を出力します。setコンテナは順序付きなので*mex.begin() で取り出すことができます（42行目）。

以下が、C＋＋プログラムとなります。

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

int main()
{
	int n, q;
	cin >> n >> q;
	vector<int> a(n);
	vector<int> s(n + 1, 0);
	for (int i = 0; i < n; ++i) {
		cin >> a[i];
		if (a[i] <= n) {
			++s[a[i]];
		}
	}

	set<int> mex;
	for (int i = 0; i <= n; ++i) {
		if (s[i] == 0) {
			mex.insert(i);
		}
	}

	for (int j = 0; j < q; ++j) {
		int i, x;
		cin >> i >> x;
		--i;
		if (a[i] <= n) {
			--s[a[i]];
			if (s[a[i]] == 0) {
				mex.insert(a[i]);
			}
		}
		a[i] = x;
		if (x <= n) {
			if (s[x] == 0) {
				mex.erase(x);
			}
			++s[x];
		}

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

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python の ordered set について

Python には、標準で順序付きの set がありません。ネットで探すといくつかの実装を見つけることができますが、記事として取り上げることを見送ります。

最後に

わたしの実力では難易度が高い問題でしたが、コンテスト中に解くことができました。setコンテナの理解が進んだ成果でしょうか。

引き続き ABC の問題を紹介していきます。