Codeforces #639 Div1 D. Résumé Review

問題文読むのしんどい…。
https://codeforces.com/contest/1344/problem/D

問題

正整数列Aが与えられる。
Aと同じ長さの数列Bのうち、以下を満たすものを求めよ。

• 0≦B[i]≦A[i]
• B[i]の総和はK
• sum(B[i]*(A[i]-B[i]^2))が最大

解法

差分を考えると、B[i]を1インクリメントしたときのsum(B[i]*(A[i]-B[i]^2))の増分は、だんだん小さくなる。
そこで、増分がx以下である回数がK回以上確保できるような最大のxを求めよう。
これはxを二分探索しつつ、B[i]を二分探索していけばよい。

int N;
ll K;
ll A[101010];
ll B[101010];
ll C[101010];

ll num(ll v) {
	int i,j;
	ll sum=0;
	FOR(i,N) {
		B[i]=0;
		for(j=30;j>=0;j--) if(B[i]+(1<<j)<=A[i]) {
			ll t=B[i]+(1<<j);
			ll d=A[i]-3*t*t+3*t-1;
			if(d>=v) B[i]+=1<<j;
		}
		
		sum+=B[i];
	}
	return sum;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) cin>>A[i];
	
	ll mi=-1LL<<62;
	for(i=62;i>=0;i--) {
		if(num(mi+(1LL<<i))>=K) mi+=1LL<<i;
	}
	/*
	
	priority_queue<pair<ll,int>> Q;
	
	ll sum=0;
	FOR(i,N) {
		D[i]=A[i]-1;
		Q.push({D[i],i});
	}
	while(K--) {
		assert(Q.size());
		x=Q.top().second;
		Q.pop();
		sum+=D[x];
		B[x]++;
		if(B[x]<A[x]) {
			D[x]=A[x]-3*B[x]*(B[x]+1)-1;
			Q.push({D[x],x});
		}
	}
	*/
	
	K=num(mi)-K;
	FOR(i,N) {
		if(K&&B[i]&&(A[i]-3*B[i]*B[i]+3*B[i]-1)==mi) {
			B[i]--;
			K--;
		}
		cout<<B[i]<<" ";
	}
	cout<<endl;
	
	
}

まとめ

ストーリー的な部分がなければ、この問題文半分以下の量になったりしないか？