torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

AtCoder Beginner Contest 178, D : Redistribution

問題概要

 正整数 $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:ごめん