Codeforces #169 Div2. E. Little Girl and Problem on Trees

さてようやく最終問題E。
http://codeforces.com/contest/276/problem/E

問題

始点を除き、辺が１つか２つしか持たない木構造が与えられる。
初期状態では各点は初期値0を持つ。
そこに対し、以下の２つのクエリを処理せよ。

• 点Vから、辺D本以内の点に対し値Xを加算する
• 点Vの値を答える。

回答

始点以外辺が高々２つということは、始点だけ複数に分岐した構造となる。
そこで、各辺毎および共通のBITを持つ。
各BITでは、始点からの距離ごとの値を持つ。

前者のクエリで値を加算するとき、点Vから辺D本の範囲に始点が入る場合、始点をはみ出る分は共通BITに加算し、それ以外の部分は辺毎のBITに加算する。
後者のクエリでは、辺毎と共通のBITの値の和を取る。

BITを作る際注意なのが、短い辺が多数ある場合と長い辺が少数ある場合があり、両方に対応できるようにBITを作ると、10^5点の要素があるBITを10^5個作るハメになりMLEする。
ここでは、要素数を最初に設定するBITを作った。

int N,Q,NP;
vector<int> E[200001];
vector<pair<int, int> > MP;
vector<vector<int> > P;

template<class V> class BIT {
public:
	int ME,BITL;
	vector<V> bit;
	
	int calcME(int v) {
		int t=1;
		while(t<v) t<<=1;
		init(t*2);
	}
	void init(int ME){
		this->ME=ME; bit.resize(ME); clear();
		BITL=0;
		while(ME>(1<<BITL)) BITL++;
	}
	void clear() {fill(bit.begin(),bit.end(),0);};
	
	// update L<=i<R
	void update(int cur,int L,int R,V v) {
		int bl=0;
		if(L==R) return;
		while(cur>=1<<(bl+1)) bl++;
		int tl=(cur-(1<<bl))<<(BITL-bl-1);
		int tr=(cur-(1<<bl)+1)<<(BITL-bl-1);
		int tm=(tl+tr)/2;
		
		if(tl==L && tr==R) {
			bit[cur]+=v;
			return;
		}
		if(L<tm) update((1<<(bl+1))+(tl>>(BITL-bl-2)),L,min(tm,R),v);
		if(R>=tm) update((1<<(bl+1))+(tm>>(BITL-bl-2)),max(L,tm),R,v);
	}
	void update(int L,int R,V v) { update(1,L,R,v);}
	
	V val(int entry) {
		V res=0;
		int i;
		FOR(i,BITL) res += bit[(1<<i)+(entry>>(BITL-i-1))];
		return res;
	}
};

BIT<int> TB[200001],MAS;

void dfs(int id, int cur, int pre) {
	MP[cur] = make_pair(id,P[id].size());
	P[id].push_back(cur);
	
	if(E[cur].size()==2) dfs(id,E[cur][0]+E[cur][1]-pre,cur);
	else TB[id].calcME(P[id].size()+3);
}

void add(int v,int x,int d) {
	if(v==0) {
		MAS.update(0,d+1,x);
	}
	else {
		pair<int, int> p = MP[v];
		if(p.second<d) {
			MAS.update(0,d-p.second,x);
			TB[p.first].update(d-p.second-1,min(p.second+d,(int)P[p.first].size())+1,x);
		}
		else {
			TB[p.first].update(p.second-d,min(p.second+d,(int)P[p.first].size())+1,x);
		}
	}
}

int get(int v) {
	if(v==0) {
		return MAS.val(0);
	}
	else {
		pair<int,int> p=MP[v];
		return TB[p.first].val(p.second)+MAS.val(p.second+1);
	}
}

void solve() {
	int f,r,i,j,k,l,x,y,z;
	ll re = 0;
	
	GET2(&N,&Q);
	FOR(i,N-1) {
		x=GETi()-1;
		y=GETi()-1;
		E[x].push_back(y);
		E[y].push_back(x);
	}
	
	NP=E[0].size();
	MP.resize(N+1);
	P.resize(NP);
	MAS.calcME(N);
	FOR(i,NP) dfs(i,E[0][i],0);
	
	FOR(i,Q) {
		if(GETi()==0) {
			x=GETi()-1;
			y=GETi();
			z=GETi();
			add(x,y,z);
		}
		else {
			_P("%d\n",get(GETi()-1));
		}
	}
	
	return;
}

まとめ

毎度BITはさらっと組めずちょくちょくバグが出る。
事前に組んでおいたライブラリもあるけど、たいてい手直しが要るし…。