Codeforces #229 Div2. E. Inna and Large Sweet Matrix

蟻本のお世話になった問題。蟻本なければ解けなかった…。
http://codeforces.com/contest/390/problem/E

問題

NxMの大きな行列がある。初期状態は全要素とも0である。
この行列に対し以下の２種類のクエリを計W個処理せよ。

1. 行列の一部長方形区間(x1,y1)-(x2,y2)と数vが与えられるので、区間に含まれる各要素にvを加える。
2. 行列の一部長方形区間(x1,y1)-(x2,y2)が与えられるので、区間内の要素の総和に対し、(P<x1またはx2<P)かつ(Q<y1またはy2<Q)となるような要素(P,Q)の総和の差を求めよ。

解法

まず後者のクエリを考える。
行列の指定区間をXとすると、要素(P,Q)に該当するのはYとなる。

YYY....YY
YYY....YY
...XXXX..
...XXXX..
...XXXX..
YYY....YY
YYY....YY
YYY....YY

これを式変形すると、以下のように(複数行の総和)＋(複数列の総和)-(全要素の総和)と一致することがわかる。

.........    ...XXXX..   YYYYYYYYY   YYY....YY
.........    ...XXXX..   YYYYYYYYY   YYY....YY
XXXXXXXXX    ...XXXX..   YYYYYYYYY   ...XXXX..
XXXXXXXXX  + ...XXXX.. - YYYYYYYYY = ...XXXX..
XXXXXXXXX    ...XXXX..   YYYYYYYYY   ...XXXX..
.........    ...XXXX..   YYYYYYYYY   YYY....YY
.........    ...XXXX..   YYYYYYYYY   YYY....YY
.........    ...XXXX..   YYYYYYYYY   YYY....YY

よって、複数行および複数列の総和と、全要素の総和を高速に計算することができればよい。

ただ、総和の計算にBit Index Treeをそのまま用いることはできない。
例えば行列中5x5の要素に3を加算する場合、各行には15ずつ加算される。
BITは単一要素の更新には対処できても、複数行に同じ値を加算する場合には対処できない。

幸い、この処理は蟻本のBITの項目に解説がある。
複数行の総和を管理する場合、BITを２つ用意する。
これまで通りの各行の総和を単純に管理するBITに加え、行番号に比例する要素を管理するBITを併用することで、このような範囲に一様に値を加算することができる。
詳細は蟻本参照。

本番は行の総和管理で２つ、列の総和管理で２つの計４つのBITを作ったが、その後このような範囲に一様に加算する処理に対応するBITライブラリを作ってみた。

template<class V, int ME> class BIT_r {
public:
	V bit[2][1<<ME];
	BIT_r(){clear();};
	void clear() {ZERO(bit);};
	
	void update(int entry, V val0, V val1) {
		entry++;
		while(entry <= 1<<ME) bit[0][entry-1]+=val0, bit[1][entry-1] += val1, entry += entry & -entry;
	}
	V total(int entry) {
		int e=entry++;
		V v0=0,v1=0;
		while(entry>0) v0+=bit[0][entry-1], v1+=bit[1][entry-1], entry -= entry & -entry;
		return e*v0+v1;
	}
	void add(int L, int R, V val) { // add val to L<=x<=R
		update(L,val,-val*(L-1));
		update(R+1,-val,val*R);
	}
};

BIT_r<ll,23> btx,bty;

int R,C,W;
ll tot;

void solve() {
	int f,i,j,k,l,x,y;
	ll x1,x2,y1,y2,v;
	
	cin>>R>>C>>W;
	
	FOR(i,W) {
		cin>>j;
		if(j==0) {
			cin>>x1>>y1>>x2>>y2>>v;
			tot += (x2-x1+1)*(y2-y1+1)*v;
			btx.add(x1,x2,v*(y2-y1+1));
			bty.add(y1,y2,v*(x2-x1+1));
		}
		else {
			cin>>x1>>y1>>x2>>y2;
			ll r1 = bty.total(y2)-bty.total(y1-1);
			ll r2 = btx.total(x2)-btx.total(x1-1);
			_P("%lld\n",r1+r2-tot);
		}
	}
	
}

まとめ

行番号比例成分を持つ、っていうのは面白いアイデア。
応用すれば、BITをN個持てばindexの(N-1)乗の項まで値を管理するデータ構造が作れそうね。