kmjp's blog

競技プログラミング参加記です

Codeforces #512 Div1 C. Putting Boxes Together

面倒な問題だね。
http://codeforces.com/contest/1053/problem/C

問題

無限に続く1次元の枠の列のうちNマスに箱が置いてあり、その位置A[i]と重さW[i]が与えられる。
以下のクエリに答えよ。

各クエリは連続する箱の区間で構成される。
箱を1動かすとコストがW[i]かかるとき、箱が連続マスに並ぶよう動かす最小コストを求めよ。
なお、他の箱を飛び越して動かすことはできない。

解法

先に位置をA'[i] = A[i]-iで座標変換しておく。
こうすると箱の大きさを気にすることなく、区間内の箱を同じ位置にまとめる問題に落とし込める。

結論から言うと、左右のバランスができるだけ等しくなる位置の箱の位置に箱を集めるとコストが最小となる。
これはある箱と隣の箱それぞれに位置に箱を集める場合のコストの差を考えると明らかである。

そこで、2つのBITを用いる。
1つめはW[i]の累積和を取るBITである。
このBITを二分探索することで、区間内の箱の重量のちょうど半分に相当する箱を求めれば、そこが箱を集めるべき位置となる。

位置がわかれば、あとはコストを求めるだけである。
これは座標*重さをmod付で累積和を取るBITをもう1つ準備しておけば高速に求められる。

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V operator()(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	void set(int e,V v) { add(e,v-val[e]);}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};

ll mo=1000000007;
template<class V, int ME> class BIT_mod {
public:
	V bit[1<<ME];
	BIT_mod(){ZERO(bit);};
	V operator()(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s%mo;}
	V add(int e,V v) { e++; while(e<=1<<ME) { bit[e-1]+=v; bit[e-1] -= (bit[e-1]>=mo)?mo:0; e+=e&-e;}}
};

BIT<ll,20> WS;
BIT_mod<ll,20> LS,RS;


int N,Q;
int A[202020],W[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	scanf("%d%d",&N,&Q);
	FOR(i,N) {
		scanf("%d",&A[i]);
		A[i]-=i;
	}
	FOR(i,N) {
		scanf("%d",&W[i]);
		WS.add(i,W[i]);
		LS.add(i,1LL*A[i]*W[i]%mo);
		RS.add(i,((1LL<<30)-A[i])*W[i]%mo);
	}
	while(Q--) {
		int L,R;
		scanf("%d%d",&L,&R);
		if(L<0) {
			L=-L;
			L--;
			WS.add(L,-W[L]);
			LS.add(L,(mo-1LL*A[L]*W[L]%mo)%mo);
			RS.add(L,(mo-((1LL<<30)-A[L])*W[L]%mo)%mo);
			W[L]=R;
			WS.add(L,W[L]);
			LS.add(L,1LL*A[L]*W[L]%mo);
			RS.add(L,((1LL<<30)-A[L])*W[L]%mo);
		}
		else {
			L--;
			R--;
			ll s=WS(R)-WS(L-1);
			ll ts=WS(L-1)+s/2;
			x=WS.lower_bound(ts);
			if(x>L) x--;
			while(x<R && WS(R)-WS(x)>=WS(x)-WS(L-1)) x++;
			int best=x;
			ll ret=RS(best-1)-RS(L-1)-(WS(best-1)-WS(L-1))%mo*((1LL<<30)-A[best]);
			ret+=LS(R)-LS(best)-(WS(R)-WS(best))%mo*(A[best]);
			cout<<(ret%mo+mo)%mo<<endl;
			
		}
		
		
	}
}

まとめ

こういう座標×重さの累積和みたいなのを、凝ったデータ構造ではなくBITみたいな軽量なデータ構造でどうにかできると気持ちよいね。