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

题目链接

题意

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

题解

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

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

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

代码

#include<bits/stdc++.h>
#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];
vector<int>e[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<n;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;
}