AtCoder ABC #257 (日鉄ソリューションズプログラミングコンテスト2022) : Ex - Dice Sum 2

なるほど…。
https://atcoder.jp/contests/abc257/tasks/abc257_h

問題

N個のサイコロがある。
それぞれの目はバラバラであり、入力としてそれぞれの値が与えられる。
また、サイコロの購入費用C[i]が与えられる。

このサイコロからK個を選び、それぞれ転がしたときに(目の総和の二乗)-(購入費用の総和)の期待値を求めよ。

解法

選んだサイコロK個について、出る目の確率変数をX[i]と置くと求める値は

で、変形すると

となる。
A[i]= E[X[i]]
B[i]=E[X[i]^2] - E[X[i]]^2 - C[i]
とすると、N個のサイコロからK個選び、(sum(A))^2+sum(B)を最大化する問題となる。

ここで、c*A[i]+B[i]が大きい順K個を選ぶという戦略を取ることを考える。
まず初期状態でc=∞とみなし、(A[i],B[i])がpair-wiseで大きい順にK個取っておくことを考える。

もしあるcが閾値c'以上では(c*A[i]+B[i])≧(c*A[j]+B[j])、c'未満では(c*A[i]+B[i])＜(c*A[j]+B[j])と反転する閾値c'があったとする。
cを徐々に小さくしていったとき、c'より小さくなった段階で、i番目よりj番目を選んだ方が得ということになる。
上記c'を事前に全(i,j)のペアに対し求めて列挙・ソートしておき、K個に含むサイコロを徐々に入れ替えながら総当たりしていこう。

本問題は解の数値を実数ではなく有理数の形で答えることが求められているので、途中計算は36倍しておくとよい。

int N,K;
int C[1010];
int A[1010][6];
ll X[1010],Y[1010];
const ll mo=998244353;
int ins[1010];

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

bool cmp(pair<pair<ll,ll>,pair<int,int>>& L,pair<pair<ll,ll>,pair<int,int>>& R) {
	ll ay=L.first.first;
	ll ax=L.first.second;
	ll by=R.first.first;
	ll bx=R.first.second;
	return ay*bx<by*ax;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) cin>>C[i];
	FOR(i,N) {
		FOR(j,6) {
			cin>>A[i][j];
			X[i]+=A[i][j];
			Y[i]+=1LL*A[i][j]*A[i][j];
		}
		Y[i]=Y[i]*6-C[i]*36-X[i]*X[i];
	}
	
	FOR(i,N) {
		FOR(j,N-1) {
			if(X[j]>X[j+1]||(X[j]==X[j+1]&&Y[j]>Y[j+1])) {
				swap(X[j],X[j+1]);
				swap(Y[j],Y[j+1]);
			}
		}
	}
	ll SX=0,SY=0;
	FOR(i,K) SX+=X[N-1-i],SY+=Y[N-1-i], ins[N-1-i]=1;
	ll ret=SX*SX+SY;
	vector<pair<pair<ll,ll>,pair<int,int>>> V;
	FOR(j,N) FOR(i,j) if(Y[i]>Y[j]) {
		V.push_back({{Y[j]-Y[i],X[j]-X[i]},{i,j}});
	}
	sort(ALL(V),cmp);
	FORR2(a,b,V) {
		x=b.first;
		y=b.second;
		if(ins[x]==0&&ins[y]==1) {
			SX+=X[x]-X[y];
			SY+=Y[x]-Y[y];
			ret=max(ret,SX*SX+SY);
			
			ins[x]=1;
			ins[y]=0;
		}
		
	}
	
	cout<<(ret%mo*modpow(36)%mo+mo)%mo<<endl;
	
	
}

まとめ

期待値の式変形の部分で戦略が良くなかったな。