AtCoder が提供しているABC（AtCoder Beginner Contest）321 のF問題をC++とPythonで解いてみました。ABC321は、2023年9月23日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
F問題 #(subset sum = K) with Add and Erase（Difficulty : 1645）
問題はリンク先をご覧ください。
ABC321 F問題 #(subset sum = K) with Add and Erase
DPの結果を戻す問題です。AtCoder Problems による Difficulty は 1645 でした。
解答案
もし、操作がタイプ1（追加）だけであれば、部分和問題と同じになります。
読みだした v に対して、i 番目の操作が終わった時に総和が j となる方法の総数 dp[i][j] には、以下の漸化式が成立します。右辺の第1項がvを選択しない場合、第2項がvを選択する場合の方法の総数となります。
- dp[i + 1][j] = dp[i][j] + dp[i][j – v] $(j \geqq v)$
- dp[i + 1][j] = dp[i][j] $(j < v) $
難しいのは、タイプ2（削除）の場合です。以下は、「けんちょんの競プロ精進記録」の記事を参考にさせていただきました。
上記の漸化式を$K + 1$ 次元の変換 $ x \rightarrow y$ と考えます。読みだした値を $v$ とすると、以下が成立します。
$$y_j = \begin{cases}
x_j + x_{j – v} & (j \geqq v) \\
x_j & (j < v)
\end{cases}$$
この $x$ と $y$ を入れ替えると以下となります。
$$x_j = \begin{cases}
y_j + y_{j – v} & (j \geqq v) \\
y_j & (j < v)
\end{cases}$$
$y_j$ について、整理すると以下になります。これが逆の操作に対応します。
$$y_j = \begin{cases}
x_j \,-\, y_{j – v} & (j \geqq v) \\
x_j & (j < v)
\end{cases}$$
この式を漸化式として表現すると以下になります。
- dp[i + 1][j] = dp[i][j] – dp[i + 1][j – v] $(j \geqq v)$
- dp[i + 1][j] = dp[i][j] $(j < v) $
C++ プログラム例（ABC321F）
漸化式をプログラムに変換しました（20ー23、28ー31行目）。この漸化式の結果は、非常に速く大きくなる場合があり、適時 998244353 の剰余を計算しています。
それぞれの操作後に dp[i + 1][k] の値を出力しています。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define MOD 998244353
int main()
{
int q, k;
cin >> q >> k;
vector<vector<ll>> dp(q + 1, vector<ll>(k + 1, 0));
dp[0][0] = 1;
for (int i = 0; i < q; ++i) {
string s;
ll x;
cin >> s >> x;
if (s == "+") {
for (int j = 0; j <= k; ++j) {
dp[i + 1][j] = dp[i][j];
if (j >= x) {
dp[i + 1][j] += dp[i][j - x];
}
dp[i + 1][j] %= MOD;
}
} else {
for (int j = 0; j <= k; ++j) {
dp[i + 1][j] = dp[i][j];
if (j >= x) {
dp[i + 1][j] += MOD - dp[i + 1][j - x];
}
dp[i + 1][j] %= MOD;
}
}
cout << dp[i + 1][k] << endl;;
}
return 0;
}上記の二次元のDPは一次元で表現することができます（in-plase 化）。タイプ2の操作の方が考えやすいです。dp[j – x] を引くために、小さな添え字から計算する必要があります（24行目）。逆に dp[j – x] を加える場合は、後ろの要素から計算する必要があります（19行目）。簡単な場合に計算すると分かりますが、前から加えていくと、1回の操作で複数の操作分の影響がでてしまいます。
この考えで書いたプログラムが以下となります。公式解説で紹介されているプログラムに近いです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define MOD 998244353
int main()
{
int q, k;
cin >> q >> k;
vector<ll> dp(k + 1, 0);
dp[0] = 1;
for (int i = 0; i < q; ++i) {
string s;
ll x;
cin >> s >> x;
if (s == "+") {
for (int j = k; j >= x; --j) {
dp[j] += dp[j - x];
dp[j] %= MOD;
}
} else {
for (int j = x; j <= k; ++j) {
dp[j] += MOD - dp[j - x];
dp[j] %= MOD;
}
}
cout << dp[k] << endl;;
}
return 0;
}どちらも AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC321F）
Python 版は、in-prace 化したプログラムを移植しました。プログラムがすっきりと書けています。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 321 F"""
MOD = 998244353
q, k = map(int, input().split())
dp = [0] * (k + 1)
dp[0] = 1
for i in range(q):
s, sx = input().split()
x = int(sx)
if s == "+":
for j in range(k, x - 1, -1):
dp[j] += dp[j - x]
dp[j] %= MOD
else:
for j in range(x, k + 1):
dp[j] += MOD - dp[j - x]
dp[j] %= MOD
print(dp[k])最後に
この問題は、コンテスト中に取り組むことすらできませんでした。コンテスト後に、いくつかの解説を読み自分なりに理解したことを紹介しました。DPの結果を戻す考え方に感心しました。
引き続き ABC の問題を紹介していきます。
