AtCoder ARC #165 : E - Random Isolation

本番2分間に合わず…。
https://atcoder.jp/contests/arc165/tasks/arc165_e

問題

N要素の木を成す無向グラフと整数Kが与えられる。
Kより大きな連結成分がある限り、そのような連結成分に対し、1頂点をランダムで選んで消すことを考える。
最終的に頂点を消す回数の期待値を求めよ。

解法

各頂点vが消される可能性を考える。

頂点vのSubTreeに対し、vを含む連結成分のサイズがa、連結成分に隣接する点(=先に消されていなければならない点)のサイズがbとする。
f(v,a,b) := 上記v,a,bを成す組み合わせ数

vを根とするとき、f(v,a,b)からvが消される確率を考えると、(a+b)要素のうち先にb要素が消され、最後にvが消される確率に等しい。
これはとなる。

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

const int NUM_=2000003;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];


unordered_map<int,ll> dp[101];
ll ret;
void dfs(int cur,int pre) {
	unordered_map<int,ll> F;
	F[1]=1;
	
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		unordered_map<int,ll> T;
		FORR2(a,b,F) {
			FORR2(a2,b2,dp[e]) {
				(T[a+a2]+=b*b2)%=mo;
			}
		}
		swap(F,T);
	}
	
	if(cur!=pre) {
		F[1000]=1;
		dp[cur]=F;
	}
	else {
		FORR2(a,b,F) {
			int aa=a/1000;
			int ab=a%1000;
			
			if(ab>K) {
				 (ret+=factr[aa+ab]*fact[ab-1]%mo*fact[aa]%mo*b)%=mo;
			}
		}
	}
	
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	inv[1]=fact[0]=factr[0]=1;
	for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
	for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	
	cin>>N>>K;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,N) {
		dfs(i,i);
	}
	cout<<ret%mo<<endl;
}

まとめ

定数倍高速化が間に合わないのもったいないな…。