AtCoder ARC 057 B

前提・実現したいこと

AtcoderのARC 057 B問題について、ほとんどの入力ファイルは通るのですが、
19.txtと56.txtでWAが出ます。
使用言語はC++です。

方針

動的計画法。
dp[tmpN][tmpW]をtmpN日ゲームをして機嫌のいい日がtmpW日となる最小の勝ち数の合計として、漸化式で計算していく。
答えはdp[N][tmpW] <= Kとなる最大のtmpW。

発生している問題・エラーメッセージ

他の同じようなWAの出方をしている人のソースコード(alpha_virginisさんのWAとAC、rpy3cppさんのWAとAC)を見ると、添え字やtmpW=1での処理周りで間違えているのだと思います。

該当のソースコード

C++
1#include <iostream>
2#include <algorithm>
3#include <cassert>
4
5using namespace std;
6
7#define MAXN 2000
8#define INF 2000000000
9
10int N, K;
11int a[MAXN + 1];
12int dp[MAXN + 1][MAXN + 1];
13
14void solve() {
15  int ans;
16  cin >> N >> K;
17  for (int i = 1; i <= N; ++i) {
18    cin >> a[i];
19  }
20
21  int sumA = 0;
22  for (int i = 1; i <= N; ++i) {
23    sumA += a[i];
24  }
25  if (sumA <= K) {
26    ans = 1;
27  }
28  else {
29    /*dp[tmpN][tmpW]は、tmpN日ゲームをして機嫌のいい日がtmpW日となる最小の勝ち数の合計*/
30    /*tmpN>=tmpWの場合だけ考える*/
31    /*tmpW=0のとき、dp[tmpN][0]=0*/
32    /*tmpN=tmpW>0のとき、tmpN=1以外は不可能*/
33    dp[1][1] = 1;
34    for (int i = 2; i <= N; ++i) {
35      dp[i][i] = INF;
36    }
37    /*tmpN>tmpW>0のとき*/
38    int A_previous_tmpN = a[1]; //前日までの勝ち数の合計
39    for (int tmpN = 2; tmpN <= N; ++tmpN) {
40      for (int tmpW = 1; tmpW <= tmpN - 1; ++tmpW) {
41        const int k_tmpN =
42          (long long int) (dp[tmpN - 1][tmpW - 1])
43          * a[tmpN] / A_previous_tmpN + 1; //tmpN日に機嫌がよくなるために必要な勝ち数
44        assert(k_tmpN > 0);
45        dp[tmpN][tmpW] = min(dp[tmpN - 1][tmpW],
46          dp[tmpN - 1][tmpW - 1] + k_tmpN);
47      }
48      A_previous_tmpN += a[tmpN];
49    }
50    for (int tmpN = 0; tmpN <= N; ++tmpN) {
51      for (int tmpW = 0; tmpW <= N; ++tmpW) {
52        assert(dp[tmpN][tmpW] >= 0);
53      }
54    }
55    for (int tmpW = N; tmpW >= 0; --tmpW) {
56      if (dp[N][tmpW] <= K) {
57        ans = tmpW;
58        break;
59      }
60    }
61  }
62
63  cout << ans << endl;
64}
65
66int main() {
67  solve();
68  return 0;
69}

補足情報

C++ (GCC 5.4.1) で提出しました。

修正したソースコード

2018/09/15作成。

C++
1#include <iostream>
2#include <algorithm>
3#include <cassert>
4
5using namespace std;
6
7#define MAXN 2000
8#define INF 1000000000 //修正
9
10int N, K;
11int a[MAXN + 1];
12int dp[MAXN + 1][MAXN + 1];
13
14void solve() {
15	int ans;
16	cin >> N >> K;
17	for (int i = 1; i <= N; ++i) {
18		cin >> a[i];
19	}
20
21	int sumA = 0;
22	for (int i = 1; i <= N; ++i) {
23		sumA += a[i];
24	}
25	if (sumA <= K) {
26		ans = 1;
27	}
28	else {
29		/*dp[tmpN][tmpW]は、tmpN日ゲームをして機嫌のいい日がtmpW日となる最小の勝ち数の合計*/
30		/*tmpN>=tmpWの場合だけ考える*/
31		/*tmpW=0のとき、dp[tmpN][0]=0*/
32		/*tmpN=tmpW>0のとき、tmpN=1以外は不可能
33		  ↑間違い a[1]=1のときはtmpN=1以外は不可能だが、a[1]>1のときは可能なケースがある*/
34		  //dp[1][1] = 1;
35		  //for (int i = 2; i <= N; ++i) {
36		  //	dp[i][i] = INF;
37		  //}
38		  /*tmpN>=tmpW>0のとき*/
39		dp[1][1] = 1;
40		int A_previous_tmpN = a[1]; //前日までの勝ち数の合計
41		for (int tmpN = 2; tmpN <= N; ++tmpN) {
42			dp[tmpN - 1][tmpN] = INF;
43			for (int tmpW = 1; tmpW <= tmpN; ++tmpW) {
44				const int k_tmpN = (a[1] == 1 && tmpW == tmpN) ? INF : (long long int) dp[tmpN - 1][tmpW - 1]
45					* a[tmpN] / A_previous_tmpN + 1;
46				//tmpN日に機嫌がよくなるために必要な勝ち数
47				assert(k_tmpN > 0);
48				dp[tmpN][tmpW] = min(dp[tmpN - 1][tmpW],
49					dp[tmpN - 1][tmpW - 1] + k_tmpN);
50			}
51			A_previous_tmpN += a[tmpN];
52		}
53		for (int tmpN = 0; tmpN <= N; ++tmpN) {
54			for (int tmpW = 0; tmpW <= N; ++tmpW) {
55				assert(dp[tmpN][tmpW] >= 0);
56			}
57		}
58		for (int tmpW = N; tmpW >= 0; --tmpW) {
59			if (dp[N][tmpW] <= K) {
60				ans = tmpW;
61				break;
62			}
63		}
64	}
65	cout << ans << endl;
66}
67
68int main() {
69	solve();
70	return 0;
71}