いろはちゃんコンテスト2019 Day2 E - 連呼(500)

E - 連呼

概要

"A"をN個、"B"をM個含むN+M文字の文字列のうち、次の条件を満たすものの通り数を求めよ。(mod 10^9+7して出力)

• 開始文字が"A"、終了文字が"B"
• 文字列中に"AAA"を含む

制約

1 ≦ N, M ≦ 10^5

方針

Comb(N, K)は2項係数を表す。
「文字列中に”AAA"を含まない」ものを考え、全体の通り数Comb(N+M-2, M-1)から引く。
このとき、文字列の形は下記の通り、「Aが1個or2個」「Bがx1(x1≧1)個」「Aが1個or2個」「Bがx2(x2≧1)個」…の繰り返しとなり、最後はBのブロックとなる。
AA BB...B A BB...B A BB..B

ブロックの個数を2*kとすると1 ≦ k ≦ min(N, M)である。
通り数はAのブロック、Bのブロックに分けて計算しかけ合わせる。
Aのブロック：N個の玉をk個の箱に1個or2個入れる⇔N-k個の玉をk個の箱に0個or1個入れる⇔Comb(k, N-k) (入れる箱を選ぶ通り数）
Bのブロック：M個の玉をk個の箱に1個以上入れる⇔M-k個の玉をk個の箱に制限なしで入れる⇔Comb(M-1, M-k) （重複組み合わせ）

Combの前処理は階乗計算+フェルマー小定理(pow(a, -1) ≡ pow(a, p-2))でO(N)。

解答

Submission #5219681 - いろはちゃんコンテスト Day2

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
 
// combination
// O(n)
 
const ll MOD = 1e9 + 7;
const ll max_n = 200003;
vector<ll> fact(max_n);
bool fact_init = false;
 
ll pow(ll a, ll b){
	if(b == 0) return 1;
	if(b % 2 == 1) return a * pow(a, b - 1) % MOD;
 
	ll d = pow(a, b / 2) % MOD;
	return d * d % MOD;
}
 
ll comb(ll n, ll k){
	if(k > n || k < 0) return 0;
	if(!fact_init){
		fact[0] = 1;
		fact[1] = 1;
 
		for(ll i = 2; i < max_n; i++){
			fact[i] = fact[i-1] * i;
			fact[i] %= MOD;
		}
 
		fact_init = true;
	}
 
	ll ret = fact[n];
	ret *= pow(fact[k],MOD-2);
	ret %= MOD;
	ret *= pow(fact[n-k],MOD-2);
	ret %= MOD;
	return ret;
 
	// X^(-1) = X^(p-2) (mod p) (Fermat's little theorem)
}
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	ll N, M;
	cin >> N >> M;
 
	ll all = comb(N + M - 2, M - 1);
 
	ll minus = 0;
	for(ll k = 1; k <= min(N, M); k++){
		minus += comb(k, N - k) * comb(M - 1, M - k) % MOD;
		minus %= MOD;
	}
 
	cout << (all - minus + MOD) % MOD << endl;
}