Codeforces #795 : Div2 F. K-Set Tree

比較的易しかった回。
https://codeforces.com/contest/1691/problem/F

問題

N頂点の木を成す無向グラフと、整数Kが与えられる。

rを、グラフを根付き木とみなしたときの根とし、Sをグラフの頂点のうちK要素を選んだものとする。
f(r,S)を、rを根としてSの各頂点をすべて含む最小の連結成分のサイズとする。
r,Sを総当たりしたときのf(r,S)の総和を求めよ。

解法

根頂点は必ずf(r,S)に含まれるので、まずその分は計上する。
各頂点vが根頂点でない場合、その頂点vがf(r,S)に計上さないケースを考える。
これは、vで区切られた各連結成分において、r及びK要素の頂点が全部1か所の連結成分内で取られたときとなるので、各連結成分のサイズがわかれば容易に計算できる。

int N,K;
vector<int> E[202020];
const ll mo=1000000007;
ll ret;

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		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;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

int dfs(int cur,int pre) {
	int V=1;
	vector<int> cand;
	FORR(e,E[cur]) if(e!=pre) {
		int C=dfs(e,cur);
		cand.push_back(C);
		V+=C;
	}
	if(N-V>0) cand.push_back(N-V);
	
	if(K==1) {
		FORR(c,cand) {
			ret+=1LL*c*(N-c)%mo;
		}
		ret+=N;
	}
	else {
		ll tot=0;
		//中心を選ぶ
		FORR(c,cand) {
			ret-=comb(N-(c+1),K-1)*c%mo*c%mo;
			tot+=comb(c,K);
		}
		
		tot=(tot%mo+mo)%mo;
		FORR(c,cand) {
			ll pat=comb(N-(c+1),K)-(tot-comb(c,K));
			ret-=pat%mo*c%mo*c%mo;
		}
		
	}
	
	return V;
	
}

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

まとめ

2750ptの割には易しい。