yukicoder : No.2369 Some Products

解法は思いついてもその後実装に手間取った。
https://yukicoder.me/problems/no/2369

問題

整数列Pが与えられる。
以下のクエリに順次答えよ。

• Pの連続部分列と正整数Kが与えられる。指定された連続部分列からK要素選んだ時の積の総和を求めよ。

解法

とする。
クエリの対象がP[L]....P[R]の時、f(L,R)のx^Kの係数が解となる。

、とする。
となるので、g(R)のi次の係数とh(L-1)の(K-i)次の係数の積を、iを総当たりして取ればよい。

g(A)、h(A)は1次式または1次式の逆数の積なので、FFTとか使わなくてもDPで計算できる。

const int mo=998244353;

int N;
vector<ll> F[5010],G[5010];
int Q;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	F[0].resize(N+1);
	G[0].resize(N+1);
	F[0][0]=G[0][0]=1;
	FOR(i,N) {
		cin>>x;
		F[i+1]=F[i];
		G[i+1]=G[i];
		for(j=1;j<=N;j++) {
			(F[i+1][j]+=F[i][j-1]*x)%=mo;
			(G[i+1][j]+=G[i+1][j-1]*(-x))%=mo;
		}
		FORR(f,F[i+1]) f=(f%mo+mo)%mo;
		FORR(f,G[i+1]) f=(f%mo+mo)%mo;
	}
	
	cin>>Q;
	while(Q--) {
		int A,B,K;
		cin>>A>>B>>K;
		ll ret=0;
		for(i=0;i<=K;i++) (ret+=F[B][i]*G[A-1][K-i])%=mo;
		cout<<ret<<endl;
	}
	
}

まとめ

最初FFT使おうとしてしまった。