Codeforces #942 : Div1 C. Fenwick Tree

そこそこ苦戦。
https://codeforces.com/contest/1967/problem/C

問題

N要素の整数列Aに対し、整数列SはS[k]がAのfenwick treeを取ったときの各値とする。
この時、関数fはS=f(A)と定義する。
整数列Bと、整数Kが与えられる。f^K(B)を答えよ。

解法

S=f(A)としたとき、Aの各要素がSのどの要素に何倍分寄与するかを求め、高速ゼータ変換の要領で足しこんでいけばよい。

int T,N;
ll K,B[202020],C[202020];
const ll 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;
}

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;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		FOR(i,N) cin>>B[i+1];
		vector<ll> cand;
		FOR(i,20) {
			cand.push_back(comb(K+i-1,i));
		}
		FOR(i,20) {
			for(j=N;j>=1;j--) if(((j)>>i<<i)==(j)&&((j)>>i)%2) {
				int cur=j;
				int step=1;
				for(x=0;x<=19;x++) if(cur<=N) {
					if(cur&(1<<(i+x))) {
						cur+=1<<(i+x);
						if(cur<=N) {
							(B[cur]+=mo-B[j]*cand[step++]%mo)%=mo;
						}
					}
				}
			}
		}
		FOR(j,N) cout<<B[j+1]<<" ";
		cout<<endl;
		
		
		
	}
	
}

まとめ

この回は問題文が簡潔で良いね。