yukicoder : No.1495 パンの仕入れ

これは他でも使えそうなテク。
https://yukicoder.me/problems/no/1495

問題

N要素で、総和がKである非負整数列Aを考える。
ここにM個の条件が与えられる。
各条件は要素xとパラメータyからなり、Aに対し(A[x]-y)^2だけコストがかかることを意味する。

Aに対し、コストの総和の最小値を求めよ。

解法

コストは下に凸な二次関数で、A[x]を1増加させたときの二階差分は常に正である。
Aのいずれかの要素を、合計Mになるまでインクリメントしていくことを考える。
当然一階差分が小さいところから順にインクリメントしていくのが良い。
Mが大きいので、愚直にインクリメントを繰り返すと間に合わないため、二分探索でM回インクリメントさせても全要素で一階差分をv以下に抑えられるようなvを求めよう。

int T;
int N,M,K;
ll A[202020],B[202020],C[202020],X[202020];

ll hoge(ll v) {
	ll sum=0;
	int i;
	FOR(i,N) {
		X[i]=0;
		if(A[i]==0) {
			if(v>=0) X[i]=K;
		}
		else {
			ll t=v-A[i]-B[i]+1;
			if(t>=0) X[i]=(t+2*A[i]-1)/(2*A[i]);
		}
		sum+=X[i];
	}
	return sum;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>M>>K;
		FOR(i,N) A[i]=B[i]=C[i]=0;
		FOR(i,M) {
			cin>>x>>y;
			x--;
			A[x]++;
			B[x]-=2*y;
			C[x]+=1LL*y*y;
		}
		
		ll ret=-1LL<<40;
		for(i=40;i>=0;i--) {
			if(hoge(ret+(1LL<<i))<=K) ret+=(1LL<<i);
		}
		hoge(ret);
		ll sumY=0;
		ll sumX=0;
		FOR(i,N) {
			sumX+=X[i];
			sumY+=A[i]*X[i]*X[i]+B[i]*X[i]+C[i];
		}
		sumY+=(ret+1)*(K-sumX);
		
		
		
		cout<<sumY<<endl;
	}
}

まとめ

本番これ出たら「M回インクリメントするの間に合わないな…」で迷っちゃいそう。