AtCoder が提供しているABC(AtCoder Beginner Contest)364 E問題をC++とPythonで解いてみました。ABC364は、2024年7月27日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Maximum Glutton(Difficulty : 1496)
問題の詳細は、リンク先をご覧ください。
DPのキーと値を入れ替えて計算量を下げます。AtCoder Problems による Difficulty は 1496 でした。
解答案
最初に上手く動かないプログラムを紹介します。
$DP_{i,j,k}$ は $i$ 個までの料理を対象として、甘さの合計 $j$ 、しょっぱさの合計 $k$ となるように選ぶ料理の最大数とします。出力する解答は、$DP_{N, j, k}$ の最大値に1を加えた値になります。プログラムは以下となります。
#include <bits/stdc++.h>
using namespace std;
#define INF 1000000001
int main()
{
int n, x, y;
cin >> n >> x >> y;
vector<vector<vector<int>>> dp(n + 1, vector<vector<int>>(x + 1, vector<int>(y + 1, -INF)));
dp[0][0][0] = 0;
for (int i = 0; i < n; ++i) {
int a, b;
cin >> a >> b;
for (int j = 0; j <= x; ++j) {
for (int k = 0; k <= y; ++k) {
dp[i + 1][j][k] = max(dp[i + 1][j][k], dp[i][j][k]);
if ((j + a <= x)&&(k + b <= y)) {
dp[i + 1][j + a][k + b] = max(dp[i + 1][j + a][k + b], dp[i][j][k] + 1);
}
}
}
}
int result = 0;
for (int j = 0; j <= x; ++j) {
for (int k = 0; k <= y; ++k) {
result = max(result, dp[n][j][k]);
}
}
result = min(result + 1, n);
cout << result << endl;
return 0;
}
構造としては、ナップザック問題と同じです。ただし、計算量は、$O(NXY) = 8 \times 10^9$ となるため、時間が間に合いません(メモリ使用量も制約を超える場合があるようです)。
$DP$ の値は、料理の最大数です。この値を甘さの合計の最小値に入れ替えます。合わせて、キーである甘さの合計は、選択した料理数に入れ替えます。つまり $DP_{i,j,k}$ を $i$ 個までの料理を対象として、$j$ 個の料理を選択した、しょっぱさの合計 $k$ となる甘さの合計の最小値とします。
この工夫により、計算量が $O(N^2Y) = 1.6 \times 10^7$ となり実行時間内に収まります。
C++ プログラム例(ABC364E)
解説したようにキー(甘さの合計値)と値(選んだ料理の最大数)を入れ替えます。dp の初期値は無限大INFとします。漸化式の表現は、プログラム20行目となります。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
#define INF 1000000001
int main()
{
int n, x, y;
cin >> n >> x >> y;
vector<vector<vector<int>>> dp(n + 1, vector<vector<int>>(n + 1, vector<int>(y + 1, INF)));
dp[0][0][0] = 0;
for (int i = 0; i < n; ++i) {
int a, b;
cin >> a >> b;
for (int j = 0; j <= i; ++j) {
for (int k = 0; k <= y; ++k) {
dp[i + 1][j][k] = min(dp[i + 1][j][k], dp[i][j][k]);
if (k + b <= y) {
dp[i + 1][j + 1][k + b] = min(dp[i + 1][j + 1][k + b], dp[i][j][k] + a);
}
}
}
}
int result = 0;
for (int j = 0; j <= n; ++j) {
for (int k = 0; k <= y; ++k) {
if (dp[n][j][k] <= x) {
result = max(result, j);
}
}
}
result = min(result + 1, n);
cout << result << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC364E)
Python版も基本的な考え方は同じです。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 364 E"""
INF = 1000000001
n, x, y = map(int, input().split())
dp = [[[INF] * (y + 1) for _ in range(n + 1)] for _ in range(n + 1)]
dp[0][0][0] = 0
for i in range(n):
a, b = map(int, input().split())
for j in range(i + 1):
for k in range(y + 1):
dp[i + 1][j][k] = min(dp[i + 1][j][k], dp[i][j][k])
if k + b <= y:
dp[i + 1][j + 1][k + b] = \
min(dp[i + 1][j + 1][k + b], dp[i][j][k] + a)
result = 0
for j in range(n + 1):
for k in range(y + 1):
if dp[n][j][k] <= x:
result = max(result, j)
result = min(result + 1, n)
print(result)
最後に
この問題では、DPのキーと値を入れ替えるという技法を学ぶことができました。
引き続き ABC の問題を紹介していきます。