CodeCraft-22 and Codeforces Round #795 (Div. 2) F. K-Set Tree

题目链接

给定一个树， f ( r , S ) f(r,S) f(r,S) 表示该树以 r r r 为根，节点 V V V 的子集 ∣ S ∣ = k |S|=k ∣S∣=k，包含 S S S 的最小子树的大小。求所有 f ( r , S ) f(r,S) f(r,S)

对每个节点计算贡献。如下图节点 4 4 4，产生的贡献为： 4 4 4 为根，其他节点选 k k k 个，且其祖先为 4 4 4；区域 A A A 的节点为根，其他节点选 k k k 个，且其祖先为 4 4 4；区域B的节点为根，其他节点选 k k k 个，且其祖先为 4 4 4…

如果以 4 4 4 为根，其他节点选 k k k 个，且祖先为 4 4 4，方案数 t o t tot tot 计算：
C ( s z [ x ] + s z [ y ] , k ) − C ( s z [ x ] , k ) − C ( s z [ y ] , k ) C(sz[x]+sz[y],k)-C(sz[x],k)-C(sz[y],k) C(sz[x]+sz[y],k)−C(sz[x],k)−C(sz[y],k)
不断对子树进行合并，上面公式的解释：两棵子树中选 k k k 个减两棵子树分别单独选 k k k 个。

而对非自身为根的方案数计算：
t o t − ( C ( n , k ) − C ( s z [ y ] , k ) − C ( n − s z [ y ] , k ) ) tot-(C(n,k)-C(sz[y],k)-C(n-sz[y],k)) tot−(C(n,k)−C(sz[y],k)−C(n−sz[y],k))
上面合并的逆

#include
#define LL long long
#define pb push_back
using namespace std;
const int N=2e5+9,mod=1e9+7;
int n,k;
LL ans;
int sz[N];
LL fac[N],inv[N];
vectore[N];
LL ksm(LL x,LL y)
{
	LL res=1;
	while(y)
	{
		if(y&1) res=res*x%mod;
		x=x*x%mod; y>>=1;
	}
	return res;
}
LL C(int x,int y)
{
	if(y>x) return 0;
	if(x==y) return 1;
	return fac[x]*inv[y]%mod*inv[x-y]%mod;
}
void dfs(int x,int fa)
{
	sz[x]=1;
	LL tot=C(1,k);
	for(int y:e[x])
	{
		if(y==fa) continue;
		dfs(y,x);
		tot=(tot+C(sz[x]+sz[y],k)-C(sz[x],k)-C(sz[y],k))%mod;
		sz[x]+=sz[y];
	}
	ans=(ans+tot*sz[x]%mod*(n-sz[x])%mod)%mod; //区域A为根的贡献
	tot=(tot+C(n,k)-C(sz[x],k)-C(n-sz[x],k))%mod;
	ans=(ans+tot*n%mod)%mod; //x为根的贡献
	for(int y:e[x])
	{
		if(y==fa) continue;
		LL res=((tot-(C(n,k)-C(sz[y],k)-C(n-sz[y],k)))%mod+mod)%mod;
		ans=(ans+res*(n-sz[y])%mod*sz[y]%mod)%mod; //孩子为根的贡献
	}
}
int main()
{
	ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	cin>>n>>k;
	for(int i=1,x,y;i
		cin>>x>>y;
		e[x].pb(y); e[y].pb(x);
	}
	fac[0]=1;
	for(int i=1;i<=n;i++)
		fac[i]=fac[i-1]*i%mod,inv[i]=ksm(fac[i],mod-2);
	dfs(1,0);
	cout<<(ans%mod+mod)%mod<<"n";
	return 0;
}