AtCoder Regular Contest 107, B : Quadruple

問題文

　https://atcoder.jp/contests/arc107/tasks/arc107_b

問題概要

　整数 $N, K$ が与えられる．整数の 4 つ組 $( a, b, c, d )$ であって，次の条件を満たすものの個数はいくつか？

• $1 \leq a, b, c, d \leq N$
• $a + b - c - d = K$

制約

• $1 \leq N \leq 10^5$
• $-2( N - 1 ) \leq K \leq 2( N - 1 )$

解法

　与式を変形すると $a + b - K = c + d$ であることが分かります．$a + b$ の値を決め打ちすることにすると，$c + d$ の値は（存在すれば）一意的に定まります（制約から，$2 \leq c + d$ です）．なので，あとは $a + b$ や $c + d$ の $a, b$ や $c, d$ の内訳としてあり得るパターンの総数をある程度高速に求められれば答えが求まります．
　公式解説によればこれは $O( 1 )$ 時間で求められますが，前計算に DP をして $\Theta( N )$ 時間かけて，一回あたり $O( 1 )$ 時間にすることもできます．具体的には，

• $\mathit{dp}_{i,j} := \text{$i$ 個の正整数で $j$ を作る方法の総数}$

として DP します．更新では $dp_{i,j}$ から $\mathit{dp}_{ i + 1, j + 1 }$ ～ $\mathit{ dp }_{ i + 1, j + N }$ へ $\mathit{ dp }_{ i, j }$ を足し込みます．これは愚直にやると時間 $\Theta( N )$ 時間かかりますが，区間への一様加算なので，いもす法を使うことで高速化できます．

コード

int main()
{
	IN( int, N, K );

	static LL dp[4][ 1 << 20 ];
	dp[0][0] = 1;
	REP( i, 2 )
	{
		REP( j, 1 << 19 )
		{
			dp[ i + 1 ][ j + 1 ] += dp[i][j];
			dp[ i + 1 ][ j + N + 1 ] -= dp[i][j];
		}
		partial_sum( ALL( dp[ i + 1 ] ), begin( dp[ i + 1 ] ) );
	}

	LL res = 0;
	REP( ab, 2, 2 * N + 1 )
	{
		const LL cd = ab - K;
		if ( 2 <= cd )
		{
			res += dp[2][ ab ] * dp[2][ cd ];
		}
	}
	cout << res << endl;

	return 0;
}