kmjp's blog

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

Codeforces #601 Div1 D. Tree Queries

Codeforcesっぽい問題?
https://codeforces.com/contest/1254/problem/D

問題

木を成す無向グラフが与えられる。
初期状態で各点には値0が設定されている。
以下のクエリを順次処理せよ。

  • 頂点vと値dが指定される。ある頂点rをランダムに選ぶ。頂点uとして、パスu-rが頂点vを通る場合、頂点uの値をd加算する。
  • 頂点vが指定される。その頂点の値の期待値を答える。

解法

平方分割で解く。
前者のタイプがD回実行されるごとに、木DPの要領で木全体の値を更新しよう。
後者のクエリが来たら、木全体の値と、未更新のD回未満の前者のクエリによる寄与分をそれぞれ足し合わせる。

int N,Q;
vector<int> E[202020];
vector<pair<int,int>> CL[202020];
int C[202020];
int L[202020],R[202020];
int id=0;
ll V[202020];
ll A[202020];

set<int> add;
const ll mo=998244353;
ll TP[202020];

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

int dfs(int cur,int pre) {
	C[cur]=1;
	L[cur]=id++;
	FORR(e,E[cur]) if(e!=pre) {
		C[cur]+=dfs(e,cur);
		CL[cur].push_back({R[e]-1,e});
	}
	R[cur]=id;
	return C[cur];
}

ll prop(int cur,int pre) {
	if(A[cur]) TP[cur]=C[cur]*A[cur];
	FORR(e,E[cur]) if(e!=pre) TP[cur]+=prop(e,cur);
	TP[cur]%=mo;
	return TP[cur];
}
void prop2(int cur,int pre,ll par) {
	TP[cur]+=par;
	if(A[cur]) TP[cur]+=mo-C[cur]*A[cur]%mo;
	(V[cur]+=TP[cur]+N*A[cur])%=mo;
	FORR(e,E[cur]) if(e!=pre) {
		ll p=TP[cur]-TP[e]+mo;
		if(A[cur]) p=(p+A[cur]*(N-C[e]))%mo;
		prop2(e,cur,p);
	}
	TP[cur]=A[cur]=0;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,-1);
	while(Q--) {
		cin>>i>>x;
		x--;
		if(i==1) {
			cin>>r;
			r=r*modpow(N)%mo;
			add.insert(x);
			(A[x]+=r)%=mo;
			
			if(add.size()>1200) {
				prop(0,-1);
				prop2(0,-1,0);
				add.clear();
			}
			
		}
		else {
			ll ret=V[x];
			FORR(a,add) {
				if(a==x) {
					(ret+=N*A[a])%=mo;
				}
				else if(L[a]<L[x] && L[x]<R[a]) {
					auto v=lower_bound(ALL(CL[a]),make_pair(L[x],0));
					(ret+=(N-C[v->second])*A[a])%=mo;
				}
				else {
					(ret+=C[a]*A[a])%=mo;
					
				}
			}
			cout<<ret<<endl;
		}
	}
}

まとめ

平方分割?に気付いてしまえば、あとはDiv1Bぐらいの実装難度の問題。