Good Bye 2022 : E. Koxia and Tree

ようやく2022年が終わる…。
https://codeforces.com/contest/1770/problem/E

問題

木を成すグラフが与えられる。
各辺には番号が振られている。

点のうちK個には蝶がいる。
ここで、各辺にそれぞれ1/2の確率で向きを設定したとする。
その後、以下の処理を行う。

• 辺の番号の順番に、辺の始点にいる蝶を一斉に終点に動かす。

2つの蝶の対における、最終的な距離の期待値の総和を求めよ。

解法

各点にいる蝶の数の期待値と、SubTree内の蝶の総和を計算しておく。
その後各辺に沿って蝶を動かすが、元々辺の両端にいる蝶は、辺の向きがどちらになっても合流する。
蝶の対に対し、辺の片側にだけ蝶がいる場合、1/2の確率でそれらは距離が1縮まる。
どっち側にもいない蝶は、辺の影響を受けない。

これを踏まえ、距離の変分を辺毎に計算していく。

int N,K;
ll B[303030],SB[303030];
vector<int> E[303030];
const ll mo=998244353;

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 U[303030],V[303030];

int id;
int L[303030],R[303030],P[303030];

void dfs(int cur,int pre) {
	P[cur]=pre;
	L[cur]=id++;
	SB[cur]=B[cur];
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		SB[cur]+=SB[e];
	}
	R[cur]=id;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,K) {
		cin>>x;
		B[x-1]=1;
	}
	FOR(i,N-1) {
		cin>>U[i]>>V[i];
		U[i]--;
		V[i]--;
		E[U[i]].push_back(V[i]);
		E[V[i]].push_back(U[i]);
	}
	dfs(0,0);
	
	ll ret=0;
	ll r2=(mo+1)/2;
	FOR(i,N-1) {
		int cur=U[i];
		int par=V[i];
		if(L[cur]<L[par]) swap(cur,par);
		assert(P[cur]==par);
		
		ll a=B[par];
		ll b=B[cur];
		
		// par->cur
		ll tot=1;
		ret+=a*(mo+1-b)%mo*r2%mo*(SB[cur]+1)%mo*(K-(SB[cur]+1))%mo;
		// cur->par
		ret+=(mo+1-a)*b%mo*r2%mo*(SB[cur]-1)%mo*(K-(SB[cur]-1))%mo;
		//それ以外
		tot-=a*(mo+1-b)%mo*r2%mo;
		tot-=b*(mo+1-a)%mo*r2%mo;
		tot=(tot%mo+mo)%mo;
		(ret+=tot*SB[cur]%mo*(K-SB[cur])%mo)%=mo;
		ret=(ret%mo+mo)%mo;
		
		
		
		
		a=(a+b)*r2%mo;
		B[par]=B[cur]=a;
	}
	ret=(ret%mo+mo)%mo;
	ll a=1LL*K*(K-1)/2%mo;
	ret=ret*modpow(a)%mo;
	
	cout<<ret<<endl;
	
}

まとめ

もうちょっとすんなり解けると良かったんだけどな。