torus711 のアレ

主に競技プログラミングの問題について書きます

Codeforces #274, Division 1, C ( Division 2, E ) : Riding in a Lift

問題概要

 n 階ある建物のエレベーターで遊ぶ。プレイヤーは最初 a 階にいて、b ( \neq a ) 階に行くことはできない。また、x 階から y 階に移動するとき、x \neq y かつ | x - y | < | x - b | でなければならない。このルールを守って k 回移動するとき、移動方法は何通りあるか、modulo で求めよ。
 n \leq 5000
 k \leq 5000

解法

 いかにも答えが大きくなるので、同一視できる状態を探して動的計画法で計算時間を削減します。移動回数は必須として、ある移動が valid か否かを判定するためには現在位置が必要になるので、次のような状態の取り方の DP が思い付きます。

  • dp[ i 階移動した ][ j 階にいる ] := 移動の総数

これだと状態数のオーダは \mathrm{ \Theta }( kn ) になりますが、各状態から次の移動を一つずつ試すと O( kn^2 ) 時間かかって TLE してしまいます。何らかの手段で更に計算時間を削る必要があります。
 ところで、現在地点 x を固定したとき、移動可能かどうかを決める不等式(問題文参照)の右辺 | x - b | は定数になります。この値を d とすると、x 階から移動できる階 y は区間 ( x - d, x + d ) 内の階となります。この区間に( Segment Tree で)暫定解を足し込むことで O( kn \log n ) 時間にできますが、k, n \leq 5000 であることを考えるとまだちょっと不安です。
 今までは所謂「配る DP」で考えてきましたが、「貰う DP」にしてみるとどうでしょうか。移動先の階 y を固定したとき、x としてありえる範囲は yb の位置関係によって決まります。y を挟んで b の反対側にある階からは、どの階からでも y に来れます。b と同じ方向の階は、yb の中間点の手前までです。ということで、貰う DP で考えた場合も遷移元となり得る階は連続した区間になっています(同じ階からの移動はできないので正確には区間は二つ)。区間の和は、予め O( n ) 時間かけて累積和を計算しておくことで O( 1 ) 時間で計算できます。累積和の計算はそれぞれの移動回数に対して一回ずつでよいので、DP 全体を O( kn ) 時間で処理できるようになります。これで \log がとれて安心できるオーダになりました。

コード

using LL = long long; using ULL = unsigned long long;
using VI = vector<int>; using VVI = vector<VI>;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define ALL( c ) (c).begin(), (c).end()

constexpr int MOD = 1000000007;

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	int n, a, b, K;
	cin >> n >> a >> b >> K;
	--a, --b;

	VVI dp( K + 1, VI( n, 0 ) );
	dp[0][a] = 1;

	REP( i, 1, K + 1 )
	{
		VI csum( 1, 0 );
		partial_sum( ALL( dp[ i - 1 ] ), back_inserter( csum ), []( const int s, const int t ){ return ( s + t ) % MOD; } );

		REP( k, 0, n )
		{
			if ( k == b )
			{
				continue;
			}

			const int d = abs( b - k );
			const int md = ( d - 1 ) / 2;

			const int l = max(     0, b < k ? k - md :     0 );
			const int r = min( n - 1, k < b ? k + md : n - 1 );

			LL s = 0;
			if ( k )
			{
				( s += csum[k] - csum[l] + MOD ) %= MOD;
			}
			if ( k + 1 < n )
			{
				( s += csum[ r + 1 ] - csum[ k + 1 ] + MOD ) %= MOD;
			}

			if ( !s )
			{
				continue;
			}
			while ( s < 0 )
			{
				s += MOD;
			}
			( dp[i][k] += s ) %= MOD;
		}
	}

	const LL res = accumulate( ALL( dp[K] ), 0LL ) % MOD;
	cout << res << endl;

	return 0;
}