kmjp's blog

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

Codeforces #232 Div1. C. On Changing Tree

これも本番割とすんなり思いつけた。
http://codeforces.com/contest/396/problem/C

問題

根付き木があり、初期状態で全頂点の値は0である。
この木に対し以下のクエリを処理せよ。

  • v,k,xが与えられるので、頂点v及びその子孫の点に対しx-(i*k)を加算する。ここでiは各頂点とvの距離である。
  • vが与えられるので、頂点vの値を返す。

解法

2種類のFenwickTreeを用いて値を管理する。
1つは各頂点における定数項pであり、もう一つは頂点からの深さの比例項qを保持する。

頂点vの根からの距離をD(v)とすると、各頂点の値はp + D(v)*qとなる。

まず各頂点をEuler Tourでたどって子孫の範囲を求めておき、加算クエリが来たら定数項にx+D(v)*kを加算し、比例項に-kを加算しておけば良い。
vの深さがD(v)ある分、定数項にD(v)*kを加算しているので最後のp + D(v)*qの計算でうまく帳尻が合う。

int N,M;
int D[300001];
int id;
int L[300001],R[300001];
vector<int> C[300001];
ll mo=1000000007;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	BIT(){clear();};
	void clear() {ZERO(bit);};

	V total(int entry) {
		V s=0;
		entry++;
		while(entry>0) {
			s+=bit[entry-1], entry -= entry & -entry;
			s%=mo;
			if(s<0) s+=mo;
		}
		return s;
	}
	void update(int entry, V val) {
		entry++;
		while(entry <= (1<<ME)) {
			bit[entry-1] += val;
			bit[entry-1] %= mo;
			if(bit[entry-1]<mo) bit[entry-1]+=mo;
			entry += entry & -entry;
		}
	}
};
BIT<ll,20> bt1,bt2;

void dfs(int cur) {
	L[cur]=id++;
	ITR(it,C[cur]) dfs(*it);
	R[cur]=id;
}

void solve() {
	int f,i,j,k,l,x,y,v;
	cin>>N;
	
	FOR(i,N-1) {
		cin>>x;
		D[i+2]=D[x]+1;
		C[x].push_back(i+2);
	}
	dfs(1);
	cin>>M;
	FOR(i,M) {
		cin>>j;
		if(j==1) {
			cin>>v>>x>>k;
			bt1.update(L[v],x+k*(ll)D[v]%mo);
			bt1.update(R[v],-(x+k*(ll)D[v]%mo));
			bt2.update(L[v],-k);
			bt2.update(R[v],k);
		}
		else {
			cin>>v;
			ll vv=bt1.total(L[v])+D[v]*bt2.total(L[v])%mo;
			vv%=mo;
			if(vv<0) vv+=mo;
			cout << vv << endl;
		}
	}
}

まとめ

なんか最近似たような処理を書いたのですんなり思いついた。
本番、1000000007でmodを取るという条件を見落として2ミスしてしまいもったいなかった。