AtCoder

ABC297 E問題(Kth Takoyaki Set)を解く

AtCoder_ABC297_E

AtCoder が提供しているABC(AtCoder Beginner Contest)297 のE問題をC++とPythonで解いてみました。ABC297は、2023年4月9日21:00に実施されました。

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

E問題 Kth Takoyaki Set(Difficulty : 1176)

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

ABC297 E問題 Kth Takoyaki Set

優先度付きキューを使って解くことができました。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 の問題を紹介していきます。

COMMENT

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

CAPTCHA