yukicoder No.389 ロジックパズルの組み合わせ

問題はこちら
No.389 ロジックパズルの組み合わせ - yukicoder

ヒントが0のみだった場合は1通り。これだけ例外処理
k-1+ΣHi>MのときNAになるのは明らか。
そうでないとき、必要な分（＝黒マスとその間の白1マス）だけ先に並べると、残った(M-(k-1+ΣHi))個の白マスを、k個ある黒カタマリに対してどの位置に配置するかk+1通りあるので、求めるべきは $_{k+1}H_{M-(k-1+\sum{H_i})}= \ _{M+1-\sum{H_i}}C_k$

#define P 1000000007
long inv[1000010];

//↓これで逆元の計算ができる！　すごい！
void makeinv(){int i;inv[1]=1;for(i=2;i<=1000000;i++)inv[i]=inv[P%i]*(P-P/i)%P;}

//コンビネーションの計算
long c(int n,int k){
	long s=1;
	int i;
	for(i=1;i<=k;i++)s=s*(n-i+1)%P*inv[i]%P;
	return s;
}

int main(){
	int n,h,k,s=0;
	makeinv();
	scanf("%d",&n);
	for(i=0;~scanf("%d",&h);k++)s+=h;
	//ΣHiがわかればいいのでHiたちを保存する必要は無い
	if(s==0)puts("1");
	else if(s+k-1>n)puts("NA");
	else printf("%lld",c(n-s+1,k));
	return 0;
}

kが与えられていないため、こんな感じで読み込むのが一番簡単だと思われる（文字列として読み込んで…とかもできるけど）

逆元の計算部分は「なんとか法」みたいな名前が付いてたような気がするんだけど、思い出せない。ご存じの方は教えて下さい
正当性は、P=ki+r に対して　i^(-1)=r^(-1)*(-k)　より直ちに明らか

とりあえずささっと縮めたい
関数化してたmakeinvとcはそのままmainの中に書くことにする
ΣHi=0のケースを場合分けするのはめんどくさいので、どうせkを求めるなら非0なものの個数をカウントすることにすれば、ΣHi=0⇔k=0となる
すると場合分けをしなくても、 $_*C_0$ なので1を返してくれる
コンビネーションは $_{M-S+1}C_k=\prod_{i=1}^{k}\frac{M-S+2-i}{i}$ を使って計算する（S=ΣHiとおいた）
もしk-1+S>MならM-S+1-k<0であり、仮定よりM-S>0なので、この積の計算途中で分子に0が混じり全体として0
逆にk-1+S<=Mなら、M<PであることからこのコンビネーションはPの倍数でない
よってこの結果が0か否かでNAの場合分けもできる

P=1e9+7;
long v[1<<20],s;n,k,i,t;
main(j){
	for(v[1]=s=scanf("%d",&n);j++<n;v[j]=v[P%j]*(P-P/j)%P);
	for(;~scanf("%d",&t);k+=!!t)n-=t;
	for(;k/++i;)s=s*(n-i+2)%P*v[i]%P;
	s=!printf(s?"%d":"NA",s);
}

で、ギュッとループ圧縮

//1つ目
for(v[1]=s=scanf("%d",&n);j++<n;v[j]=v[P%j]*(P-P/j)%P);for(;~scanf("%d",&t);k+=!!t)n-=t;
for(v[1]=s=scanf("%d",&n);j++<n?v[j]=v[P%j]*(P-P/j)%P:~scanf("%d",&t)?n-=t,k+=!!t:0;);

//2つ目
for(v[1]=s=scanf("%d",&n);j++<n?v[j]=v[P%j]*(P-P/j)%P:~scanf("%d",&t)?n-=t,k+=!!t:0;);for(;k/++i;s=s*v[i]%P*(n-i+2)%P);
for(v[1]=s=scanf("%d",&n);j++<n?v[j]=v[P%j]*(P-P/j)%P:~scanf("%d",&t)?n-=t,k+=!!t:k/++i?s=s*v[i]%P*(n-i+2)%P:0;);

完成したものがこちら

long v[1<<20];
P=1e9+7,s,n,k,i,t;
main(j){
	for(v[1]=s=scanf("%d",&n);j++<n?v[j]=v[P%j]*(P-P/j)%P:~scanf("%d",&t)?n-=t,k+=!!t:k/++i?s=s*v[i]%P*(n-i+2)%P:0;);
	s=!printf(s?"%d":"NA",s);
}

179B
vはlongで無くてもいいけど、そうすると掛け算の箇所で1L*を2回書く必要があり、トータルで1B損
逆元の計算をせずにpowで押し切る方法も考えたけど、短くなりそうにはなかった

メモ

yukicoderでゆるふわgolf

yukicoder No.389 ロジックパズルの組み合わせ