ABC312 D問題（Count Bracket Sequences）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）312 のD問題をC++とPythonで解いてみました。ABC312は、2023年7月29日21:00に実施されました。

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

D問題 Count Bracket Sequences（Difficulty : 964）

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

ABC312 D問題 Count Bracket Sequences

DPを使って解くことができます。AtCoder Problems による Difficulty は 964 でした。

解答案

DP（Dynamic Programming: 動的計画法）を使います。問題の見た目は異なりますが、漸化式の処理が ABC291D問題（解説）と近い問題でした。

dp[i][j] は、i 文字目までにある ‘?’ を ‘(‘ または ‘)’ に置き換える方法の個数とします。j は、そこまでの文字に含まれる「'(‘ の数 ー ‘)’ の数」とします。ただし、文字列のどの時点でもこの数は負になることはないとします（このため j は 0 以上の場合だけを考えます）。これは、’)’ の数が ‘(‘ より多くならないことを示します。

• s[i] が ‘(‘ か ‘?’ の場合
  • dp[i + 1][j + 1] に dp[i][j] を加える。
    これは、'(‘ の数がひとつ増えることに該当します。
• s[i] が ‘)’ か ‘?’ の場合（ただし、j – 1 は負にならないとする）
  • dp[i + 1][j – 1] に dp[i][j] を加える。
    これは、’)’ の数がひとつ増えて、「'(‘ の数 ー ‘)’ の数」が一つ減ることを意味します。

漸化式の初期値は、以下となります。

• dp[0][0] = 1、それ以外の dp[i][j] = 0 とする。

解答は、dp[N][0] となります。制約のもとで計算結果は非常に大きくなるため、適宜 998244353 との割った余りを求めていきます。

C++ プログラム例（ABC312D）

上記の漸化式をプログラムとして表現しました（15、18ー26行目）。

以下が、C＋＋プログラムとなります。

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long int ull;

#define MOD 998244353

int main()
{
	string s;
	cin >> s;
	int n = s.length();

	vector<vector<ull>> dp(n + 1, vector<ull>(n + 1, 0));
	dp[0][0] = 1;
	for (int i = 0; i < n; ++i) {
		for (int j = 0; j < n; ++j) {
			if ((s[i] == '(')||(s[i] == '?')) {
				dp[i + 1][j + 1] += dp[i][j];
				dp[i + 1][j + 1] %= MOD;
			}
			if (((s[i] == ')')||(s[i] == '?'))
			 && (j - 1 >= 0)) {
				dp[i + 1][j - 1] += dp[i][j];
				dp[i + 1][j - 1] %= MOD;
			}
		}
	}

	cout << dp[n][0] << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC312D）

Python 版は C++ 版を移植しました。紹介したプログラムは、Python（3.8.2）と PyPy3 （7.3.0）で実行時間の差が大きくでました。Python では、27テストケース中、17テストケースで TLE 判定（実行制限時間: 2sec）となりました。PyPy3 では、最長のテストケースでも 374ms しか時間を消費せず AC 判定となりました。

"""AtCoder Beginner Contest 312 D"""
MOD = 998244353

s = input()
n = len(s)

dp = [[0 for i2 in range(n + 1)] for i1 in range(n + 1)]
dp[0][0] = 1
for i in range(n):
    for j in range(n):
        if s[i] == '(' or s[i] == '?':
            dp[i + 1][j + 1] += dp[i][j]
            dp[i + 1][j + 1] %= MOD
        if (s[i] == ')' or s[i] == '?') and j - 1 >= 0:
            dp[i + 1][j - 1] += dp[i][j]
            dp[i + 1][j - 1] %= MOD

print(dp[n][0])

最後に

解説を読み、過去に似た問題（ABC291D問題）を解いていたことを思い出しました。まだ DP の習熟度が足りないようです。のんびりと長く構えて学んでいくことにします。

引き続き ABC の問題を紹介していきます。