kmjp's blog

競技プログラミング参加記です

AtCoder ABC #285 : Ex - Avoid Square Number

ARC

そんなテクもあったな。
https://atcoder.jp/contests/abc285/tasks/abc285_h

問題

整数N,Kと長さKの整数列Eが与えられる。
Eは、素数のうち小さい順にK個を使った、ある整数の素因数分解の結果を示すものとする。

積が上記整数となるような、N要素の正整数列を考える。
全要素が平方数でないようなものは何通りか。

解法

各要素、少なくとも1つの素因数で位数が奇数ならよい。
包除原理の要領で考える。
f(i) := N要素中i要素が平方数であり、残りi要素は平方数かどうかわからない
とすると、求める値は
\displaystyle \sum_{i=0}^N (-1)^{N-i} Binom(N,i) f(i)
である。

\displaystyle f(i) = \prod_{j=1}^K [x^{E_j}] (1+x^2+x^4+...)^i(1+x+x^2+x^3+...)^{N-i}
だが、右辺の無限級数の部分は
\displaystyle \frac{1}{(1+x)^j(1-x)^N}
と書ける。1/(1-x)^Nは、結局累積和をN回取ること、(1+x)^jは、符号を正負入れ替えながら累積和をj回取ること、と考えると、max(E)+1要素の数列を2N回累積和取れば、この分数の部分を計算できる。

int N,K;
int E[101010];
const ll mo=1000000007;
int dp[10101][10101];

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>>K;
	FOR(i,K) cin>>E[i];
	
	dp[N][0]=1;
	for(i=N-1;i>=0;i--) {
		FOR(x,10001) {
			dp[i][x]=dp[i+1][x];
			if(x) (dp[i][x]+=dp[i][x-1])%=mo;
		}
	}
	ll ret=1;
	FOR(j,K) (ret*=dp[0][E[j]])%=mo;
	for(i=1;i<=N;i++) {
		FOR(x,10001) {
			dp[i][x]=dp[i-1][x];
			if(x) (dp[i][x]+=mo-dp[i][x-1])%=mo;
		}
		ll pat=comb(N,i);
		FOR(j,K) (pat*=dp[i][E[j]])%=mo;
		if(i%2==0) {
			ret+=pat;
		}
		else {
			ret+=mo-pat;
		}
	}
	cout<<ret%mo<<endl;
}

まとめ

あえてNTTを使わせない制約ってのがいいね。