AtCoder Beginner Contest 178, D : Redistribution

問題文

　https://atcoder.jp/contests/abc178/tasks/abc178_d

問題概要

　正整数 $S$ が与えられる．数列であって，すべての項が $3$ 以上かつ総和が $S$ になるものの個数を $\pmod{ 10^9 + 7 }$ で求めよ．

制約

• $1 \leq S \leq 2{,}000$

解法

　次のように状態をとる動的計画法を考えます．

• dp[ 数列の長さ ][ 総和 ] := # of ways

ここで，dp[i][j] からの遷移は，$3$ 以上 $2{,}000$ 以下の k について

• dp[i][j] を dp[ i + 1 ][ j + k ] に足し込む

です．これは $O( S^3 )$ 時間のアルゴリズムで，単純に実装すると TLE します．一方で，dp[i][j] = 0 な状態からの遷移を continue 等で無くすと謎の理由で 間に合います*1．

コード

constexpr int MOD = 1'000'000'007;

int main()
{
	IN( int, S );

	static int dp[ 1024 ][ 2048 ];
	dp[0][0] = 1;

	REP( i, 667 )
	{
		REP( j, 2001 )
		{
			if ( !dp[i][j] )
			{
				continue;
			}

			REP( k, 3, 2001 )
			{
				if ( S < j + k )
				{
					continue;
				}

				( dp[i][ j + k ] += dp[i][j] ) %= MOD;
			}
		}
	}

	int res = 0;
	REP( i, 668 )
	{
		( res += dp[i][S] ) %= MOD;
	}
	cout << res << endl;

	return 0;
}

*1:ごめん