TopCoder SRM 744 Div1 Easy, Div2 Hard ModularQuadrant

Easy落としたけどMedium早解きでレートキープ。
https://community.topcoder.com/stat?c=problem_statement&pm=15236

問題

無限に広がる二次元グリッドがある。
セル(r,c)には数値max(r,c)%3が書かれているものとする。
グリッド内の矩形区間(r1,c1)-(r2,c2)が指定されるので、その中の数値の総和を求めよ。

解法

f(r,c) := (0,0)-(r,c)の区間の数値の総和
とすると解はf(r2,c2)-f(r1-1,c2)-f(r2,c1-1)+f(r1-1,c1-1)になるので、あとはf(r,c)を求めればよい。

r≧cとすると、正方形区間(0,0)-(c,c)と長方形区間(c+1,c)-(r,c)を別々に考えるとわかりやすい。
内部の数値の分布がいずれも単純なのでO(1)で計算できる。

class ModularQuadrant {
	public:
	ll sq(ll a) {
		ll ret=0;
		
		if(a%3==1) ret+=(2*a+1), a--;
		if(a%3==0) a--;
		if(a>0) ret+=(a+1)*(a+1)+4*(a+1)/3;
		return ret;
		
	}
	ll hoge(ll a,ll b) {
		if(a<0 || b<0) return 0;
		if(a<b) swap(a,b);
		ll ret=0;
		
		while(a>b && (a-b)%3) ret+=(a%3)*1LL*(b+1), a--;
		ret+=(a-b)*(b+1);
		return ret+sq(b);
	}
	
	long long sum(int r1, int r2, int c1, int c2) {
		return hoge(r2,c2)-hoge(r2,c1-1)-hoge(r1-1,c2)+hoge(r1-1,c1-1);
	}
	
}

まとめ

本番はO(√N)でゴリ押そうとしたけど、関係ないところでミスして落ちた。