AtCoder Regular Contest 036 C: 偶然ジェネレータ

問題

C: 偶然ジェネレータ - AtCoder Regular Contest 036 | AtCoder

解法

dp[pos][diff0][diff1]:= pos文字目まで見た時に、(0の文字数)-(1の文字数)の最大値がdiff0であり、(1の文字数)-(0の文字数)の最大値がdiff1である組み合わせの数

とする。

S[pos]='0' のとき [pos][diff0][diff1] から [pos][diff0+1][diff1-1] に遷移し、 S[pos]='1' のとき [pos][diff0][diff1] から [pos][diff0-1][diff1+1] に遷移する。
S[pos]='?' のときは両方に遷移する。

コード

import java.util.Scanner;

class Main {
	private static final int MOD = 1000000007;

	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int N = sc.nextInt();
		int K = sc.nextInt();
		char[] S = sc.next().toCharArray();
		sc.close();

		/*
		 * dp[pos][diff0][diff1]:= pos文字目まで見た時に、 (0の文字数)-(1の文字数)の最大値がdiff0であり、
		 * (1の文字数)-(0の文字数)の最大値がdiff1である組み合わせの数
		 */
		int[][][] dp = new int[N + 1][K + 2][K + 2];
		dp[0][0][0] = 1;

		for (int pos = 0; pos < N; pos++) {
			for (int diff0 = 0; diff0 <= K; diff0++) {
				for (int diff1 = 0; diff1 <= K; diff1++) {
					if (dp[pos][diff0][diff1] == 0) {
						continue;
					}

					if (S[pos] == '0' || S[pos] == '?') {
						if (diff1 > 0) {
							dp[pos + 1][diff0 + 1][diff1 - 1] += dp[pos][diff0][diff1];
							dp[pos + 1][diff0 + 1][diff1 - 1] %= MOD;
						} else {
							dp[pos + 1][diff0 + 1][0] += dp[pos][diff0][0];
							dp[pos + 1][diff0 + 1][0] %= MOD;
						}
					}
					if (S[pos] == '1' || S[pos] == '?') {
						if (diff0 > 0) {
							dp[pos + 1][diff0 - 1][diff1 + 1] += dp[pos][diff0][diff1];
							dp[pos + 1][diff0 - 1][diff1 + 1] %= MOD;
						} else {
							dp[pos + 1][0][diff1 + 1] += dp[pos][0][diff1];
							dp[pos + 1][0][diff1 + 1] %= MOD;
						}
					}
				}
			}
		}

		int ans = 0;
		for (int i = 0; i <= K; i++) {
			for (int j = 0; j <= K; j++) {
				ans += dp[N][i][j];
				ans %= MOD;
			}
		}

		ans %= MOD;
		System.out.println(ans);
	}
}

感想

「最大値」を引数にもつのがポイントな気がする。次出てきた時解けるのか……？