AtCoder が提供しているABC(AtCoder Beginner Contest)294 のD問題をC++とPythonで解いてみました。ABC294は、2023年3月19日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Bank(Difficulty : 385)
問題はリンク先をご覧ください。
コンテナを使い、クエリを処理する問題です。AtCoder Problems による Difficulty は 385 でした。
解答案
コンテストに提出したプログラムを紹介します。
- 並んでいる人を priority_queue コンテナ waited に格納します。
- 呼ばれた人を別の priority_queue コンテナ called に格納します。
- C++ の priority queue には削除がないため、受付された人は、set コンテナ receipt に格納します。called の top が receipt にないことを確認して表示します。
用意したコンテナに対して、クエリの処理結果を格納していきます。以下がC++プログラムになります。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, q;
cin >> n >> q;
priority_queue<int, vector<int>, greater<int>> waited;
priority_queue<int, vector<int>, greater<int>> called;
set<int> receipt;
for (int i = 1; i <= n; ++i) {
waited.push(i);
}
for (int i = 0; i < q; ++i) {
int command;
cin >> command;
if (command == 1) {
called.push(waited.top());
waited.pop();
} else if (command == 2) {
int x;
cin >> x;
receipt.insert(x);
} else if (command == 3) {
while (true) {
int temp = called.top();
if (receipt.find(temp) != receipt.end()) {
called.pop();
} else {
break;
}
}
cout << called.top() << endl;
}
}
return 0;
}
C++ プログラム例(ABC294D)
上記は、冗長なプログラムになっていました。
- waited は必ず番号順になるため、コンテナは必要なく変数で管理します。
- called は、削除ができる set コンテナに変更します。
set コンテナは順番を保持しており、最初の要素が一番小さくなります。結果的に以下のすっきりとしたプログラムになりました。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, q;
cin >> n >> q;
set<int> called;
int waited = 1;
for (int i = 0; i < q; ++i) {
int command;
cin >> command;
if (command == 1) {
called.insert(waited);
++waited;
} else if (command == 2) {
int x;
cin >> x;
called.erase(x);
} else if (command == 3) {
cout << *(called.begin()) << endl;
}
}
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC294D)
Python では、辞書を使います。辞書は、Python 3.7 から順序を保つようになっています。辞書の要素は、pop で削除して(12行目)、next(iter()) で先頭要素を取得できます(14行目)。
なお、このプログラムは、Pyhon(3.8.2)では、TLE(Time Limit Exceeded)判定でした。PyPy3(7.3.0)は、AC判定でした。
"""AtCoder Beginner Contest 294 D"""
n, q = map(int, input().split())
called = {}
waited = 1
for _ in range(q):
que = list(map(int, input().split()))
if que[0] == 1:
called[waited] = 1
waited += 1
elif que[0] == 2:
called.pop(que[1])
else:
print(next(iter(called)))
最後に
問題の意図に従ったプログラムをコンテストで提出しました。公式回答を元にコンテナの使い方を見直し、すっきりとしたプログラムになりました。
引き続き ABC の問題を紹介していきます。