AtCoder が提供しているABC(AtCoder Beginner Contest)310 のE問題をC++とPythonで解いてみました。ABC310は、2023年7月15日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 NAND repeatedly(Difficulty : 1261)
問題はリンク先をご覧ください。
DPを使って解くことができました。AtCoder Problems による Difficulty は 1261 でした。
解答案
文字列の先頭から、結果が 0 になる場合と 1 になる場合をカウントします。
dp[j][d] を $i \leqq j$ の範囲で $f(i, j) = d$ となる個数とします。
s[i] が 0 の場合、$i < j$ の範囲の漸化式は以下となります。
- dp[i][0] = 0
s[i] が 0 であるため、結果は 0 になることはありません。 - dp[i][1] = dp[i – 1][0] + dp[i – 1][1]
s[i] が 0 であるため、ひとつ前の結果に関わらず結果は 1 となります。
s[i] が 1 の場合、$i < j$ の範囲の漸化式は以下となります。
- dp[i][0] = dp[i – 1][1]
その前の演算結果が 1 の場合、反転して 0 になります。 - dp[i][1] = dp[i – 1][0]
その前の演算結果が 0 の場合、反転して 1 になります。
これに $i = j$ の場合を加えます。つまり s[j] が 0 か 1 によって、dp[i][0 or 1] の値をインクリメントします。
初期値は以下となります。
- dp[0][0] = 1、dp[0][1] = 0 (s[0] が 0の場合)
- dp[0][1] = 0、dp[0][1] = 1 (s[0] が 1の場合)
C++ プログラム例(ABC310E)
上記の考察に従い、漸化式を実装します。求める結果は、$(0 \leqq j < N)$ の範囲の dp[j][1] の和となります(26ー30行目)。
漸化式の処理コードの背景色を変更しています。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
string s;
cin >> s;
vector<vector<ll>> dp(n, vector<ll>(2, 0));
++dp[0][s[0] - '0'];
for (int i = 1; i < n; ++i) {
if (s[i] - '0' == 0) {
dp[i][0] = 0;
dp[i][1] = dp[i - 1][0] + dp[i - 1][1];
} else {
dp[i][0] = dp[i - 1][1];
dp[i][1] = dp[i - 1][0];
}
++dp[i][s[i] - '0'];
}
ll result = 0;
for (int i = 0; i < n; ++i) {
result += dp[i][1];
}
cout << result << endl;
return 0;
}
紹介したプログラムでは、dp を二次元の変数にしました。最初の添え字を無くして、dp を更新することができます。その代わり変数 result を都度更新する必要があります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
string s;
cin >> s;
ll result = 0;
vector<ll> dp(2, 0);
++dp[s[0] - '0'];
result += dp[1];
for (int i = 1; i < n; ++i) {
if (s[i] - '0' == 0) {
dp[1] = dp[0] + dp[1];
dp[0] = 0;
} else {
swap(dp[0], dp[1]);
}
++dp[s[i] - '0'];
result += dp[1];
}
cout << result << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC310E)
C++ の2つ目のプログラムを移植しました。多重代入ができる(11、13行目)ため、すっきりとしたプログラムになりました。
"""AtCoder Beginner Contest 310 E"""
n = int(input())
s = input()
result = 0
dp = [0] * 2
dp[int(s[0])] += 1
result += dp[1]
for i in range(1, n):
if int(s[i]) == 0:
dp[0], dp[1] = 0, dp[0] + dp[1]
else:
dp[0], dp[1] = dp[1], dp[0]
dp[int(s[i])] += 1
result += dp[1]
print(result)
こちらも「AC」と判定されました。
最後に
この問題は、コンテスト中に解くことができませんでした。DP で解けるというヒントから実装できました。
公式解説では、添え字を減らしたバージョンに近いプログラムを紹介していました。ただし、自分で、これをいきなり書くことは難しいと感じました。まずは多少効率が悪くても思いつくことができるプログラム(2次元版)をコンテストで書けるようなりたいです。
引き続き ABC の問題を紹介していきます。