Codeforces #564 Div1 C. Nauuo and Pictures

これHardの時は思いつかなかったな…。
https://codeforces.com/contest/1172/problem/C2

問題

N種類の絵があり、初期状態で各価値はW[i]である。
M日間の間、美術館にいずれかの絵が飾られるが、その確率は価値に比例した確率となる。

なお、各絵には館主の好き嫌いが設定されている。
好きな絵は1回飾られると価値が1増え、嫌いな絵は1減る。

各絵について、飾られる回数の期待値を求めよ。

解法

EasyはN,Mの値が小さい。
よって、各絵について、(対象の絵を飾った回数, 他に好きな絵を飾った回数, 他に嫌いな絵を飾った回数)を状態として、その状態に至る確率をDPで求めると、O(NM^3)で解ける。

HardはN,Mの値が大きい。
ここで重要な知見は、好きな絵・嫌いな絵それぞれ、飾られる回数は価値に比例するということである。
とすると、価値1の好きな絵・嫌いな絵が飾られる回数を求めればよい。
これも上記のとおりに状態を取るとまだO(M^3)かかるが、以下の通りO(M^2)のDPを2回行えば済むようになる。

• 自分と同じ属性の絵がx回飾られる場合、初期状態で価値1の絵が飾られる回数の期待値
• 好きな絵を飾った回数と嫌いな絵を飾った回数に対するその状態に至る確率

int N,M;
int A[202020];
ll W[202020];
ll SW[2];
ll mo=998244353;

ll dp[3030][3030];

inline int mulmod(int a,int b,int mo=mo) {
	int d,r;
	if(a==0 || b==0) return 0;
	if(a==1 || b==1) return max(a,b);
	__asm__("mull %4;"
	        "divl %2"
		: "=d" (r), "=a" (d)
		: "r" (mo), "a" (a), "d" (b));
	return r;
}

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

ll comb(int P_,int Q_) {
	static const int N_=4020;
	static ll C_[N_][N_];
	
	if(C_[0][0]==0) {
		int i,j;
		FOR(i,N_) C_[i][0]=C_[i][i]=1;
		for(i=1;i<N_;i++) for(j=1;j<i;j++) C_[i][j]=(C_[i-1][j-1]+C_[i-1][j])%mo;
	}
	if(P_<0 || P_>N_ || Q_<0 || Q_>P_) return 0;
	return C_[P_][Q_];
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	FOR(i,N) cin>>A[i];
	
	FOR(i,N) {
		cin>>W[i];
		SW[A[i]]+=W[i];
	}
	
	dp[0][0]=1;
	FOR(x,M) FOR(y,M) if(x+y<M && dp[x][y]) {
		ll a=SW[1]+x;
		ll b=SW[0]-y;
		ll v=dp[x][y]*modpow(a+b)%mo;
		(dp[x+1][y]+=v*a)%=mo;
		(dp[x][y+1]+=v*b)%=mo;
	}
	
	ll S[2]={};
	ll Xs=1,As=1;
	int p,q;
	for(p=0;p<=M;p++) {
		ll sum=0,Ys=1,As2=As;
		for(q=0;p+q<=M;q++) {
			ll a=Ys*modpow(As2)%mo;
			(sum+=a*dp[p+q][M-(p+q)]%mo*comb(p+q,p))%=mo;
			Ys=mulmod(Ys,SW[1]-1+q);
			As2=mulmod(As2,SW[1]+p+q);
		}
		(Xs*=1+p)%=mo;
		(S[1]+=Xs*sum)%=mo;
		(As*=(SW[1]+p))%=mo;
	}
	Xs=As=1;
	ll Ys=1;
	for(q=0;q<=min(M,(int)SW[0]-1);q++) {
		ll a=Ys*modpow(As)%mo;
		(S[0]+=a*dp[M-q][q])%=mo;
		(Ys*=(SW[0]-1-q))%=mo;
		(As*=(SW[0]-q))%=mo;
	}
	
	FOR(i,N) cout<<S[A[i]]*W[i]%mo<<endl;
	
}

まとめ

初期の重みに比例する、ってのに気付けるのすごいな。