Codeforces #1053 : Div1. D. Attraction Theory

そこまでコードは長くない。
https://codeforces.com/contest/2150/problem/D

問題

N要素の整数列Aが与えられる。
初期状態で、N要素の整数列PはP[i]=iを満たす。

ここで、Pに対し以下の処理を任意回数行える。

• 1以上N以下の正整数Xを指定する。
  • Pのうち、X未満の値を持つ要素はインクリメントする
  • Pのうち、Xより大きい値を持つ要素はデクリメントする

こうしてできるP全パターンに対し、sum(A[P[i]])の総和を求めよ。

解法

F[j] := 最終的にP[i]=jとなるP[i]の個数
とすると、Fのうち正の値を取る部分だけ抽出すると、その間はすべて正であり、かつ両端以外は奇数しかとらない。
そこで両端の偶奇だけ総当たりすると、残りの値の埋め方は二項係数で計算できる。

int T,N;
ll A[202020];
ll S[202020];
const ll mo=998244353;
const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll comb(ll N_, ll C_) {
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(ll P_,ll Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	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;
	
	cin>>T;
	while(T--) {
		cin>>N;
		FOR(i,N) {
			cin>>A[i];
			S[i+1]=(S[i]+A[i])%mo;
		}
		ll ret=S[N]*N%mo;
		ll Lsum=0,Rsum=0,Tsum=0;
		for(int len=N;len>=2;len--) {
			(Lsum+=A[N-len])%=mo;
			(Rsum+=A[len-1])%=mo;
			if(len==N) {
				Tsum=S[N];
			}
			else {
				Tsum+=S[N]-S[N-len]+mo;
				Tsum+=mo-(S[N]-S[len]);
				Tsum%=mo;
			}
			for(x=1;x<=2;x++) for(y=1;y<=2;y++) {
				int s=N-x-y-(len-2);
				if(s<0||s%2) continue;
				ll pat=hcomb(len,s/2);
				(ret+=pat*Tsum)%=mo;
				if(x==2) (ret+=pat*Lsum)%=mo;
				if(y==2) (ret+=pat*Rsum)%=mo;
				(ret+=pat*s%mo*Tsum%mo*inv[len])%=mo;
			}
		}
		cout<<ret%mo<<endl;
	}
}

まとめ

言われてみると難しくないんだよな。