kmjp's blog

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

yukicoder : No.1549 [Cherry 2nd Tune] BANning Tuple

yukicoder

これも実装が面倒。
https://yukicoder.me/problems/no/1549

問題


N要素の非負整数列Xを考える。
条件iは、以下を定める。

  • X[K[i]]は、A[i]以上B[i]以下である。

以下のQ個のクエリに答えよ。

  • クエリjは、2値S[j]、T[j]で指定される。
  • 条件1～jを満たすXのうち、総和がS[j]以上T[j]以下となるのは何通りか。

解法

以下の多項式を考える。
P(x,i) := A(x,i,n)*x^n (A(x,i,n)は、X[i]が値nを取れるなら1、取れないなら0)

各クエリは、条件1～jを満たすP(x,n)に対し、
Q(x) = Prod(P(x,1),...,P(x,N))
を求めて、Q(x)のS[j]次～T[j]次の係数の和を取ると解となる。

最大値ではNが大きくQは小さいので、ほとんどのXは条件の影響を受けない。
そこで、K[i]に登場しないXの添え字nについて、重複数え上げの要領でP(x,n)を求めておこう。

SegTreeの要領でProd(***)の部分を計算していくことを考える。
条件が1個増えると、P(logQ)個だけ多項式の積を更新する必要があるが、NTTを使えば何とか間に合う。

ll N;
int Q;
ll K[101];
int A[101],B[101],S[101],T[101];

const int mo=998244353;
ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}


template<class T> vector<T> fft(vector<T> v, bool rev=false) {
	int n=v.size(),i,j,m;
	
	for(i=0,j=1;j<n-1;j++) {
		for(int k=n>>1;k>(i^=k);k>>=1);
		if(i>j) swap(v[i],v[j]);
	}
	for(int m=2; m<=n; m*=2) {
		T wn=modpow(5,(mo-1)/m);
		if(rev) wn=modpow(wn);
		for(i=0;i<n;i+=m) {
			T w=1;
			for(int j1=i,j2=i+m/2;j2<i+m;j1++,j2++) {
				T t1=v[j1],t2=(ll)w*v[j2]%mo;
				v[j1]=t1+t2;
				v[j2]=t1+mo-t2;
				while(v[j1]>=mo) v[j1]-=mo;
				while(v[j2]>=mo) v[j2]-=mo;
				w=(ll)w*wn%mo;
			}
		}
	}
	if(rev) {
		ll rv = modpow(n);
		FOR(i,n) v[i]=(ll)v[i]*rv%mo;
	}
	return v;
}

template<class T> vector<T> MultPoly(vector<T> P,vector<T> Q,bool resize=false) {
	if(resize) {
		int maxind=0,pi=0,qi=0,i;
		int s=2;
		FOR(i,P.size()) if(norm(P[i])) pi=i;
		FOR(i,Q.size()) if(norm(Q[i])) qi=i;
		maxind=pi+qi+1;
		while(s*2<maxind) s*=2;
		P.resize(s*2);Q.resize(s*2);
		if(s<=16) { //fastpath
			vector<T> R(s*2);
			for(int x=0;x<2*s;x++) for(int y=0;x+y<2*s;y++) (R[x+y]+=P[x]*Q[y])%=mo;
			return R;
		}
	}
	P=fft(P), Q=fft(Q);
	for(int i=0;i<P.size();i++) P[i]=(ll)P[i]*Q[i]%mo;
	return fft(P,true);
}

ll comb(ll P_,ll Q_) {
	if(P_<0 || Q_<0 || Q_>P_) return 0;
	ll p=1,q=1;
	Q_=min(Q_,P_-Q_);
	P_%=mo;
	for(int i=1;i<=Q_;i++) p=p*P_%mo, q=q*i%mo,P_--;
	return p*modpow(q,mo-2)%mo;
}

ll hcomb(ll P_,ll Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

vector<ll> seg[256];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	
	for(j=128;j<256;j++) {
		FOR(x,4096) seg[j].push_back(1);
		FOR(x,4096) seg[j].push_back(0);
	}
	
	
	map<ll,int> M;
	FOR(i,Q) {
		cin>>K[i]>>A[i]>>B[i]>>S[i]>>T[i];
		if(M.count(K[i])==0) {
			x=M.size();
			M[K[i]]=x;
		}
		K[i]=M[K[i]];
	}
	for(j=M.size();j<128;j++) for(x=1;x<4096;x++) seg[128+j][x]=0;
	FOR(x,4096) seg[128+M.size()][x]=hcomb(N-M.size(),x);
	for(j=127;j>=1;j--) {
		seg[j]=MultPoly(seg[2*j],seg[2*j+1]);
		seg[j].resize(4096);
		seg[j].resize(8192);
	}
	FOR(i,Q) {
		for(x=A[i];x<=B[i];x++) seg[128+K[i]][x]=0;
		j=128+K[i];
		while(j>1) {
			j/=2;
			seg[j]=MultPoly(seg[2*j],seg[2*j+1]);
			seg[j].resize(4096);
			seg[j].resize(8192);
		}
		ll ret=0;
		for(x=S[i];x<=T[i];x++) ret+=seg[j][x];
		cout<<ret%mo<<endl;
		
	}
}

まとめ

Codeforcesだと要素数10^5とかのFFTをよく使うけど、最近AtCoderだと小さな多項式を何個も掛け合わせるの目にするよね。