torus711 のアレ

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

HHKB プログラミングコンテスト 2020, D : Squares

問題概要

 2 次元空間に直交座標系をとる.ここに,辺長が $N$ で辺が軸並行な白い正方形を置く.
 更に,辺が軸並行で辺長がそれぞれ $A, B$ で,色がそれぞれ青と赤の正方形を,各頂点が格子点に載るように配置することを考える.このとき,青と赤の正方形は白い正方形の領域に内包されなければならず,また,青と赤の正方形は重なってはならない(共通部分の面積は $0$ でなければならない).そのような配置の総数を(MOD で)求めよ.
 ……という問題に,一つの入力ファイルにつき $T$ 回答えよ.

制約

  • $1 \leq T \leq 10^5$
  • $1 \leq N \leq 10^9$
  • $1 \leq A, B \leq N$

解法

 入力ファイルあたり $T$ 個もの問題を解かなければならず,また,各問題毎に入力が独立にくるので,前計算等は(たぶん)できません.よって,各問題を $O( \log N )$ 時間ぐらいで解かないと間に合いません.その前提で,各計算を $O( 1 )$ 時間ぐらいで行う方針を考えていきます.
 求めるべき答えを言い換えると,青と赤の正方形を重なりに関する条件を無視して配置する方法の総数から,必ず重なるように配置する方法の総数を引いたものと言うことができます.青い正方形を重なりの条件を無視して配置する方法の総数は,縦の自由度と横の自由度がそれぞれ $N - A + 1$ である*1ことから,$( N - A + 1 )^2$ です.同様に,赤い正方形を重なりの条件を無視して配置する方法は $( N - B + 1 )^2$ です.
 ここで 2 乗の形になることの解釈は,正方形を $X$ 軸・$Y$ 軸にそれぞれ射影してできる線分の方を決めているからだと思うことができます.逆に,射影される線分が決まれば,正方形も一意的に定まります.このことを踏まえると,軸ごとに独立に考えてもよいことになります.実際,長さ $N$ の線分に長さ $A$ の線分を載せる方法の総数というのは $N - A + 1$ 通りになることが確かめられ(,$B$ についても同様で),軸ごとにバラしたものを二乗しながら足せばよいことが分かります.
 必ず重複するような配置の仕方について考えますが,これ以降は一次元版で考えて最後に二乗します.つまり,青と赤の線分が必ず重なるように長さ $N$ の線分上に配置する方法の総数を返す関数を $\mathit{ overlap }$ と呼ぶことにすれば,求めるべき答えは $( N - A + 1 ) ^ 2 + ( N - B + 1 ) ^ 2 + \mathit{ overlap }( N, A, B )^2$ で,$\mathit{ overlap }$ の中身について一次元版で考えます.
 青と赤の線分が重なりをもつ場合というのは,次の 2 パターンに分けられます.

  • 青と赤の線分の内,片方が他方を完全に内包する場合
  • 青と赤の線分が部分的に重なりをもつ場合

 前者の場合は比較的簡単で,青と赤の線分両方によって張られる線分の長さは $\max( A, B )$ になります.この線分を長さ $N$ の線分に乗せる方法の総数は,$N - \max( A, B ) + 1$ 通りです.また,短くない方の線分の内部にある,長くない方の線分の配置の方法が,同様の計算によって $\max( A, B ) - \min( A, B ) + 1$ 通りであることも分かり,この積が前者のケースの置き方の総数です.
 後者の場合はやや複雑になります.まず,青と赤の線分両方によって張られる線分の長さは,$\min( N, A + B - 1 )$ です($N$ を超えてはいけないこと・必ず長さ $1$ だけ青と赤で重複しなければならないこと,などを反映するとこの形になります).重ね方も 1 通りではなく,重ねた後の長さが $\max( A, B ) + 1$ になるものから $\min( N, A + B - 1 )$ になるもののすべてを考えなければなりません.しかしよく考えると,長さ $N$ の線分に載せる線分の長さを $1, 2, 3, \dots, N$ と動かしていくとき,置き方の総数はそれぞれ $N, N - 1, N - 2, \dots, 1$ です.ここから,短すぎる線分の分(先程の条件に合わない分)を除けばよいです.長さ $\min( N, A + B - 1 )$ 未満の線分を配置する方法は $1$ から $N - \min( N, A + B - 1 )$ の和で,長さ $\max( A, B )$ 以下の部分交差ではない配置の方法は $1$ から $N - \max( A, B + 1 )$ の和です.$0$ から $x$ の和は $\frac { x ( x + 1 ) } 2$ と $O( 1 )$ 時間で求められます.また,部分交差の場合は青を左に出すか赤を左に出すかで 2 パターンあるので,実際には「割る $2$」は打ち消されます.
 これで先程の場合分けの前者と後者の両方が求まったことになり,必要な値が揃いました.各計算は $O( 1 )$ 時間で計算できているので,入力ファイル全体では $O( T )$ 時間となります.

コード

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

LL int_sum( const LL N )
{
	assert( N < MOD );
	return N * ( N + 1 ) / 2 % MOD;
}

LL int_sum( const LL L, const LL R )
{
	assert( L < MOD );
	assert( R < MOD );
	assert( L <= R );
	return ( int_sum( R ) - int_sum( L - 1 ) ) % MOD;
}

LL solve_olverlap( const LL N, const LL A, const LL B )
{
	const LL max_w = min( N, A + B - 1 );
	const LL min_w = max( A, B );
	LL res = 0;
	if ( min_w < max_w )
	{
		( res += int_sum( N - max_w + 1, N - min_w ) % MOD ) %= MOD;
		( res *= 2 ) %= MOD;
	}
	( res += ( N - min_w + 1 ) * ( max( A, B ) - min( A, B ) + 1 ) % MOD ) %= MOD;
	return res;
}

LL solve()
{
	IN( LL, N, A, B );
	LL res = 1;
	const LL a = N - A + 1;
	const LL b = N - B + 1;
	const LL overlap = solve_olverlap( N, A, B );
	( res *= a * a % MOD ) %= MOD;
	( res *= b * b % MOD ) %= MOD;
	( res += MOD - overlap * overlap % MOD ) %= MOD;
	return res;
}

int main()
{
	IN( int, T );
	REP( T )
	{
		cout << solve() << '\n';
	}
	cout << flush;

	return 0;
}

*1:小さい例を脳内で動かすと分かります