AtCoder が提供しているABC（AtCoder Beginner Contest）297 のE問題をC++とPythonで解いてみました。ABC297は、2023年4月9日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Kth Takoyaki Set（Difficulty : 1176）
問題はリンク先をご覧ください。
優先度付きキューを使って解くことができました。AtCoder Problems による Difficulty は 1176 でした。
解答案
$N$ の制約が、$1 \leqq N \leqq 10$ と少ない値に設定されています。仮にこの $N$ 種類のたこやきの値段を、a[0]、a[1]、…、a[9] とします。
一番安い金額は、a[0]、a[1]、…、a[9] の一番安い値段になります。仮に a[0] が安かったとします。次に安い金額の候補は、a[1]、…、a[9] に加えて、a[0] + a[0]、a[0] + a[1]、…、a[0] + a[9] になります。
一番少ない数は、優先度付きキューによって取り出すことができます。取り出した値に対して、a[0]、a[1]、…、a[9] を加えた値をキューにプッシュします。
キューに積まれる数の最大値は、$K \times N$ となります。1回のアクセスの計算量は、$O(\log(N K))$ となります。プッシュする回数も考慮すると $N K$ 回アクセスします。合わせて、計算量は、$O(N K \log(N K))$ となります。
注意点は、問題の制約から「同じ金額を支払う方法が複数存在する場合は1 回だけ数えます」とあります。
C++ プログラム例（ABC297E）
上記の手順をプログラムにします。
C++ の優先度付きキュー（priority_queue）は、デフォルトで最大値を取り出すため、最小値を取り出すよう宣言します（15行目）。一番安い値段の候補として、a[0] から a[n-1] をキューにプッシュします（16-18行目）。
金額は配列 result に格納していきます。同じ金額か確認するために、result[index] とキューから取り出した値を比較しています（28行目）。値が異なる場合に result を更新して（29、30行目）、新しい候補をキューにプッシュします（34-36行目）。
result[0] には、0 を格納しているため、K 番目の安い価格は、result[k] になります（39行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
ull n, k;
cin >> n >> k;
vector<ull> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
priority_queue<ull, vector<ull>, greater<ull>> que;
for (int i = 0; i < n; ++i) {
que.push(a[i]);
}
vector<ull> result(k + 1);
result[0] = 0;
int index = 0;
while (index < k) {
ull next;
while (!que.empty()) {
next = que.top();
que.pop();
if (next != result[index]) {
++index;
result[index] = next;
break;
}
}
for (int i = 0; i < n; ++i) {
que.push(next + a[i]);
}
}
cout << result[k] << endl;
return 0;
}2024/5/2追記
setを使ったプログラムも紹介します。setを使う場合、重複を調べる必要がなくなります。priority_queue を使う場合より時間はかかるようです。以下が、setバージョンです。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
ull n, k;
cin >> n >> k;
vector<ull> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
set<ull> que;
for (int i = 0; i < n; ++i) {
que.insert(a[i]);
}
vector<ull> result(k + 1);
result[0] = 0;
int index = 0;
while (index < k) {
ull next = *(que.begin());
que.erase(next);
++index;
result[index] = next;
for (int i = 0; i < n; ++i) {
que.insert(next + a[i]);
}
}
cout << result[k] << endl;
return 0;
}どちらも AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC297E）
Python の優先度付きキューは、heapq モジュールをインポートして使います。heapq の使い方は、日本語の公式解説を参考にしました。
C++版との違いは以下の2点です。
- 優先度付きキュー heapq は、最小の要素を返します。
- heapq.heappop で、値を返し、pop を行います。
- que が空か確認していません。空の場合は、IndexError が送出されます。この問題では、que が空になることはありません。
- 変数 next の名前が組込み関数と重複しているため、next_price に変更しました。
"""AtCoder Beginner Contest 297 E"""
import heapq
n, k = map(int, input().split())
a = list(map(int, input().split()))
que = []
heapq.heapify(que)
for i in range(n):
heapq.heappush(que, a[i])
result = [0] * (k + 1)
result[0] = 0
index = 0
while index < k:
while True:
next_price = heapq.heappop(que)
if next_price != result[index]:
index += 1
result[index] = next_price
break
for i in range(n):
heapq.heappush(que, next_price + a[i])
print(result[k])こちらも「AC」と判定されました。
最後に
この問題は、N重ループを考えて、それは現実的ではないことに気づいて、それ以上進みませんでした。優先度付きキューについては、理解していたつもりですが、典型的な場面で使うことができませんでした。とはいえ、これで学べたと思います。
引き続き ABC の問題を紹介していきます。
