paizaラーニングレベルアップ問題集の「十億連勝」をやってみた

paiza×Qiita記事投稿キャンペーンということで、paizaラーニングレベルアップ問題集の[十億連勝](https://paiza.jp/works/mondai/real_event/continuous_winning)をやってみました。
- - -
この問題は、各ステージに入る時点で、
- 既に$X$連勝達成したか、未達成か（$X$連勝を超えた場合は、ボーナスステージは発生しないので無視）
- 前のステージで全勝したか、途中で負けてしまったか
- 前のステージで全勝した場合、トータル何連勝中か（あるいは第何ステージから連勝が続いているか）

そして、
- 当該ステージで$X$連勝を超えることがあるか

を考慮する必要があります。

そこで、$S_i(f,w)$（$0\le i\le N$, $f\in \\{\\text{false},\text{true}\\}$, $0\le w\le X$）を、
- $S_i(\text{false},w)$：第$i$ステージ終了時点で最大連勝数$X$未満で、現在$w$連勝中となる通り数。但し、$S_i(\text{false},X)$も正の値を持ちうるものとします。
- $S_i(\text{true},w)$：第$i$ステージ終了時点で最大連勝数$X$を達成済みで、現在$w$連勝中となる通り数。

とします。$S_0(\text{false},0)=1$, $S_0(\text{false},w)=0$（$0<w\le X$）, $S_0(\text{true},w)=0$（$0\le w\le X$）で初期化すると、
- $w+A_i\le X$の（第$i$ステージで全勝しても連勝数は$X$を超えない）場合
    - 第$i$ステージで全勝した場合：$S_{i+1}(f,w+A_i)=S_i(f,w)$
    - 途中で負けてしまった場合：第$i$ステージの勝数として$0,1,...,A_i-1$の$A_i$通りが考えられるので、$S_{i+1}(f,0)$に$A_i\times S_i(f,w)$を足し合わせる。
- $w+A_i>X$の（第$i$ステージで全勝すると$X$連勝を超える）場合
    - 第$i$ステージで$X-w$勝を超えた場合はボーナスステージが発生しなくなるので、考えない。
    - 第$i$ステージで丁度$X-w$勝して次に負けた場合：$S_{i+1}(\text{true},0)$に$S_i(\text{true},w)$と$S_i(\text{false},w)$を足し合わせる。
    - 第$i$ステージで$X-w$勝未満で負けた場合：第$i$ステージの勝数として$0,1,...,X-w-1$の$X-w$通りが考えられるので、$S_{i+1}(f,0)$に$(X-w)\times S_i(f,w)$を足し合わせる。

最終的に、$S_N(\text{false},X)+\sum_wS_N(\text{true},w)$が答えとなります。

さて、$S_i(f,w)$の管理についてかんがえてみましょう。すると
```
states[N][2][X] // 2番目の「2」はfalse, trueを管理する
```
のような三次元配列が考えられますが、今回の制約条件は$X\le10^9$なので、こんな巨大な配列を用意することはできません。但し、第$i$ステージまでの連勝数は第何ステージから連勝が続いているかによるので、高々$N$の制約条件である$4000$種類です。そこで、
```
states[N][2] : 連想配列<連勝数, 通り数>
```
の様に持たせておきましょう。更に、$S_{i+1}(f,w)$について、計算に必要なのは$S_i(f,w)$だけです。つまり、古い$S_i(f,w)$については保持しておく必要がないので
```
states[2] : 連想配列<連勝数, 通り数>
new_states[2] : 連想配列<連勝数, 通り数>
```
の2変数を用意しておきます。
今回は、$1,000,000,000$で割った余りを求めますので、計算の途中でも$1,000,000,000$で割った余りにしながら答えを求めていきます。

以下、C++のサンプルコードです。
```cpp
#include <iostream>
#include <vector>
#include <map>
#include <utility>
using namespace std;

const long MOD = 1000000000;

int main() {
	int n;
	long x;
	cin >> n >> x;
	vector<long> A(n);
	for (int i = 0; i < n; i++)
		cin >> A[i];
	vector<map<long, long>> states(2);
	states[0][0] = 1L;
	for (int i = 0; i < n; i++) {
		vector<map<long, long>> temp(2);
		for (int j = 0; j < 2; j++) {
			for (pair<long, long> state : states[j]) {
				long wins = state.first;
				if (wins + A[i] <= x) {
					temp[j][wins + A[i]] = (temp[j][wins + A[i]] + states[j][wins]) % MOD;
					temp[j][0] = (temp[j][0] + A[i] * states[j][wins]) % MOD;
				} else {
					temp[1][0] = (temp[1][0] + states[j][wins]) % MOD;
					temp[j][0] = (temp[j][0] + (x - wins) * states[j][wins]) % MOD;
				}
			}
		}
		states = temp;
	}

	long answer = states[0][x];
	for (pair<long, long> state : states[1]) {
		answer = (answer + state.second) % MOD;
	}
	cout << answer << endl;
}
```
- - -
連想配列を持たない言語でも、第$i$ステージまでに何ステージ連勝が続いているかということと、[累積和](https://paiza.jp/works/mondai/prefix_sum_problems)を用いることにより
```
states[2][N]
new_states[2][N]
```
と、二次元配列を2個用意することにより解くことができます。

以下、C言語のサンプルプログラムです。
```c
#include <stdio.h>
#include <string.h>

const long MOD = 1000000000;

int main() {
	int n;
	long x;
	scanf("%d %ld", &n, &x);
	long A[n];
	long S[n + 1]; // 勝利数累積和
	S[0] = 0L;
	for (int i = 0; i < n; i++) {
		scanf("%ld", &A[i]);
		S[i + 1] = S[i] + A[i];
	}
	long states[2][n + 1];
	memset(states, 0, sizeof(states));
	states[0][0] = 1L;
	for (int i = 0; i < n; i++) {
		long temp[2][n + 1];
		memset(temp, 0, sizeof(temp));
		for (int j = 0; j < 2; j++) {
			for (int k = 0; k <= i; k++) {
				long wins = S[i] - S[i - k]; // 直近kステージ連勝した時の連勝数
				if (wins > x) { // x連勝超えは考える必要がないのでbreak
					break;
				}
				if (wins + A[i] <= x) {
					temp[j][k + 1] = (temp[j][k + 1] + states[j][k]) % MOD;
					temp[j][0] = (temp[j][0] + A[i] * states[j][k]) % MOD;
				} else {
					temp[1][0] = (temp[1][0] + states[j][k]) % MOD;
					temp[j][0] = (temp[j][0] + (x - wins) * states[j][k]) % MOD;
				}
			}
		}
		memcpy(states, temp, sizeof(temp));
	}
	long answer = 0L;
	for (int i = 0; i <= n; i++) {
		if (S[n] - S[n - i] > x) { // x連勝超えは考える必要がない
			break;
        }
		answer = (answer + states[1][i]) % MOD;
		if (S[n] - S[n - i] == x) { // 丁度x連勝で第nステージ終了
			answer = (answer + states[0][i]) % MOD;
        }
	}
	printf("%ld\n", answer);
	return 0;
}
```