AtCoder AGC #021 : E - Ball Eat Chameleons

またややこしい設定。
https://agc021.contest.atcoder.jp/tasks/agc021_e

問題

N匹の青いカメレオンがいる。
ここにK個の赤か青のボールを投げこむ。

ボールを投げるといずれかのカメレオンがそれを食べる。
カメレオンが最後に食べたボールにより、赤青のボール数の差が生じると、カメレオンがその色になる。

K個投げ込んだ結果、カメレオンがすべて赤くなるようなボールの色順は何通りか。

解法

ボール数が等しいときは直前の状態を引き継ぐので少しややこしい。
カメレオンが赤いためには、食べた赤ボールが青ボールより多いか、同数だが最後に食べたボールが青（言い換えれば最後の直前は赤を多く食べていた状態である）ことが求められる。

これを踏まえ、赤ボール数Rを総当たりすることを考えよう。
青ボール数をBとすると、R≧Bでなければならないのは明らかである。

R＞Bの場合

最小で(N-(R-B))体は先に赤を1個たべ、あとに1個青を食べることになる。
すると、残り(R-B)体が食べるのはR-(N-(R-B))個の赤とB-(N-(R-B))=R-N個の青であり、これは前者の方が大きいので途中経過は問わず最後は赤くなる。
そこで、(N-(R-B))体が「赤を1個たべ、あとに1個青を食べる」ケースを考える。
投げるボール順を見て、赤が多い間は問題ないが、青が多くなると(N-(R-B))体の方は赤くなれないものが生じかねない。
ただしB-(N-(R-B))=R-N個までなら残り1体が食べる分として吸収できる。

これを踏まえると、この組み合わせは以下のようになる。
２次元座標で(0,0)から(R,B)までxまたはy座標が増加する向きに格子点を辿る組み合わせを考える。
ただしy-x＞R-N、すなわりy-x≧R-N+1となるケースは認められない。
これはカタラン数の求め方を応用すると容易に求められ、となる。

R==Bの場合

全てのカメレオンが赤青同数のボールを食べる。
よって少なくとも最後の１体が食べる最後のボールは赤くなければならない。
これは(0,0)から(R,B)ではなく(0,0)から(R,B-1)に至る経路を数え上げればよい。

int N,K;
ll mo=998244353;

ll combi(ll N_, ll C_) {
	const int NUM_=1400001;
	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;
}

ll catalan_arrange(int X,int Y,int T) {
	return (combi(X+Y,Y)-combi(X+Y,Y-T)+mo)%mo;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	
	ll ret=0;
	for(int R=N;R<=K;R++) {
		int B=K-R;
		if(R==B) {
			ret+=catalan_arrange(R,B-1,(R-N+1));
		}
		if(R>B) {
			ret+=catalan_arrange(R,B,(R-N+1));
		}
	}
	
	cout<<ret%mo<<endl;
}

まとめ

前のCFでも出たけど、カタラン数の数え上げ問題苦手。