yukicoder : No.1608 Yet Another Ants Problem

これはよいアレンジ問題。
https://yukicoder.me/problems/no/1608

問題

長さLの棒に、N匹のアリがおり、その座標が与えられる。
ある時間からアリがそれぞれ棒の両端どちらかの向きに一斉に動き始める。
アリは同じ速さで移動し、もし他の蟻とぶつかると、互いにその時点から逆向きに動き始める。
棒の端に到達したアリは棒から落下する。

N匹の蟻が左右どちらに動くかは2^N通り考えられる。
それぞれ、N匹の各アリが最後に落ちるケースは何通りか数えよ。

解法

全アリが同じ向きに動く例は自明。

アリが衝突せずすれ違うと考えると、最後に落ちるアリは右向きに移動する最も左のアリか、左向きに移動する最も右のアリのいずれかである。
よって、何番目のアリがそれぞれ上記ケースに該当するか、O(N^2)パターン数えることを考える。

この時、例えば左からL番目のアリが右向きに移動する最も左のアリで、左からR番目のアリが左向きに移動する最も右のアリとする。
そして、この時前者のアリが最後に落下するとする。
衝突を考慮すると、L番とR番の間に左向きに移動するアリがx匹いた場合、最後に落下するのはL+x番目のアリである。
またその組み合わせはComb(R-L+1,x)通りである。

xを総当たりするとO(N^3)かかり間に合わないので高速化することを考える。
i匹めのアリが最後に落ちる組み合わせ数を格納する配列をP[i]とする。
上記総当たりの中で行う処理は、P[L]～P[R]に、Comb(R-L+1,x)を加算する処理といえる。

これはパスカルの三角形の要領で、数え上げることが可能。
最初のアリのパターンの列挙がO(N^2)、その後パスカルの三角形の計算がO(N^2)で全体をO(N^2)で解くことができる。

int N,T;
int A[3030];
const ll mo=998244353;
ll pat[3030];
ll co[3030][3030];

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>T;
	T=2*T-1;
	FOR(i,N) {
		cin>>A[i];
		A[i]*=2;
	}
	
	co[0][0]++; // all right
	co[0][N-1]++; // all left
	FOR(i,N-1) {
		if(T-A[i]<A[i+1]) co[0][i]++;
		else co[0][i+1]++;
	}
	for(int R=0;R<N;R++) {
		for(int L=R+1;L<N;L++) {
			int ri=N-L;
			int le=R+1;
			if(T-A[R]<A[L]) {
				co[N-(ri+le)][le-1]++;
			}
			else {
				co[N-(ri+le)][le]++;
			}
		}
	}
	
	for(i=N-1;i>=0;i--) {
		FOR(j,N) {
			co[i][j]+=co[i+1][j];
			if(j) co[i][j]+=co[i+1][j-1];
			co[i][j]%=mo;
		}
	}
	
	FOR(i,N) cout<<co[0][i]<<endl;
	
}

まとめ

割と手間取った。