TCO '10 Round3 Hard Passwords

最近は毎日、書くまでもないくらいの難易度の問題を解いてます。
更新量少なくなってしまっているので記事にできるくらいの、もう少し難しめの問題を解きたいorz

問題

英小文字大文字、数字からなるパスワードを考える。
パスワードのうち、長さがNで、小文字を最低L文字含み、大文字を最低U文字含み数字をD文字含むようなものはいくつあるか、mod 10^9+9で求めよ。

制約条件

N≦200000

方針

Nがそこそこ大きいので、
dp[i][j][k]:i文字で、小文字をj文字使い、大文字をk文字使うフレーズの数のようなDPは無理。


Editorialにあったのはこんな解法。


まず、大文字小文字を同一視した場合の数を考えてみる。
これは、Aをアルファベットの文字数として、
L+U < A <= N-DなるAについて、C(N,A)*26^A*10^(N-A)を足し合わせればよい。


次にA個のアルファベットのうち、大文字がU文字以上、小文字がL文字以上のパターンは何通りかを考える。
これは01からなるN桁の列で、0の数がL以上、1の数がU以上のものの数に相当する。
これについては以下のような漸化式が成り立つ


cnt[A]=2*cnt[A-1] + C(A-1,L-1) + C(A-1,U-1)


なぜならA桁のパターンというのは、

  • A-1桁のパターンに任意に0,1を付け足す
  • A-1桁で、0がLに1個だけ足りないパターンに0を付け足す
  • A-1桁で、1がUに1個だけ足りないパターンに1を付け足す

ものの和に等しいからである。


階乗および階乗の逆元を求めておくことで、C(n,r)はO(1)で計算できる。
累乗もあらかじめ計算しておく。

ソースコード

const int mod=1000000009;
ll pw10[200001],pw26[200001];
ll f[200001],invf[200001];
ll cnt[200001];

ll extgcd(ll a, ll b, ll &x, ll &y) {
  ll g = a; x = 1; y = 0;
  if (b != 0) g = extgcd(b, a % b, y, x), y -= (a / b) * x;
  return g;
}
ll invMod(ll a, ll m) {
  ll x, y;
  if (extgcd(a, m, x, y) == 1) return (x + m) % m;
  else                         return 0; // unsolvable
}
ll C(int n,int r){
	return f[n]*invf[r]%mod*invf[n-r]%mod;
}

class Passwords {
public:
  int countValid(int N, int L, int U, int D) {
  	
  	pw10[0]=pw26[0]=1;
  	rep(i,N){
  		pw10[i+1]=(pw10[i]*10)%mod;
  		pw26[i+1]=(pw26[i]*26)%mod;
  		cnt[i+1]=0;
  	}
  	
  	f[0]=1; invf[0]=invMod(1,mod);
  	rep(i,N){
  		f[i+1]=f[i]*(i+1)%mod;
  		invf[i+1]=invMod(f[i+1],mod);
  	}
  	
  	cnt[0]=L+U==0?1:0;
  	for(int A=max(1,L+U);A<=N;A++){
  		cnt[A]=2*cnt[A-1];
  		if(L>0)cnt[A]+=C(A-1,L-1);
  		if(U>0)cnt[A]+=C(A-1,U-1);
  		cnt[A]%=mod;
  	}
  	
  	ll ans=0;
  	for(int A=L+U;A<=N-D;A++){
  		ll a=cnt[A];
  		a=a*pw26[A]%mod;
  		a=a*pw10[N-A]%mod;
  		a=a*C(N,A)%mod;
  		ans=(ans+a)%mod;
  	}
  	return ans;
  }
};