CodeCraft-19 and Codeforces Round #537 (Div. 2) E 虚树 + 树形dp(新坑)

https://codeforces.com/contest/1111/problem/E

题意

一颗有n个点的树,有q个询问,每次从树挑出k个点,问将这k个点分成m组,需要保证在同一组中不存在一个点是另一个点的祖先这种情况,问有多少中分组方案

题解

• 首先解决转移问题,假设一次询问全颗树的分组方案数,定义dp[u][i]为到树节点u,分了i组的方案数,其中num(fa[u])为u的祖先节点个数
• \(dp[u][i]=dp[fa][i]*(i-num(fa[u]))+dp[fa][i-1]\)
  1. 当前这个点独立分组:\(dp[fa][i-1]->dp[u][i]\)
  2. 当前这个点加入前面的分组:\(dp[fa][i]*(i-num(fa[u]))->dp[u][i]\)
• 反着扫,滚掉一维
• 然后每次询问一些点,但点询问的总规模数和总点数相同的这种问题可以考虑用虚树解决

代码

#include<bits/stdc++.h>
#define MOD 1000000007
#define MAXN 100005
#define MAXM 305
#define MAXK 100005
#define ll long long
using namespace std;
int n,i,q,u,v,sp[MAXN],dfn[MAXN],dep[MAXN],f[MAXN][40],rt;
ll ans,dp[MAXM];
int k,m,r,ti,S[MAXK];
vector<int>g[MAXN];
vector<int>vt[MAXK];

bool cmp(int x,int y){
	return dfn[x]<dfn[y];
}
void dfs1(int u,int fa){
	dfn[u]=++ti;
	dep[u]=dep[fa]+1;
	f[u][0]=fa;
	for(int i=1;i<=20;i++)f[u][i]=f[f[u][i-1]][i-1];
	for(auto v:g[u])
		if(v!=fa)dfs1(v,u);
}

int lca(int u,int v){
	if(dep[u]<dep[v])swap(u,v);
	int dis=dep[u]-dep[v];
	for(int i=20;i>=0;i--)if((dis>>i)&1)u=f[u][i];
	if(u==v)return u;
	for(int i=20;i>=0;i--){
		if(f[u][i]!=f[v][i]){
			u=f[u][i];
			v=f[v][i];
		}
	}
	return f[u][0];
}
void add_edge(int u,int v){
	vt[u].push_back(v);
	vt[v].push_back(u);
}
void build_vtree(vector<int>& a){
	sort(a.begin(),a.end(),cmp);
	a.erase(unique(a.begin(),a.end()),a.end());
	int t=0;
	S[t++]=a[0];
	vector<int>b;b.clear();
	for(int i=1;i<a.size();i++){
		if(t==0){S[t++]=a[i];continue;}
		int l=lca(a[i],S[t-1]);
		while(t>1&&dfn[S[t-2]]>=dfn[l]){add_edge(S[t-2],S[t-1]);t--;}
		if(l!=S[t-1]){add_edge(S[t-1],l);S[t-1]=l;b.push_back(l);}
		S[t++]=a[i];
	}
	while(t>1){add_edge(S[t-1],S[t-2]);t--;}
	for(auto u:b)a.push_back(u);
}
void dfs2(int u,int fa,int cnt){
	if(sp[u]){
		for(int i=m;i>=1;i--){
			if(i>cnt)dp[i]=dp[i]*(i-cnt)%MOD;
			else dp[i]=0;
			dp[i]=(dp[i]+dp[i-1])%MOD;
		}
		dp[0]=0;
		cnt++;
	}
	for(int i=0;i<vt[u].size();i++){
		int v=vt[u][i];if(v==fa)continue;
		dfs2(v,u,cnt);
	}
}

int main(){
	cin>>n>>q;
	for(int i=0;i<n-1;i++){
		scanf("%d%d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs1(1,0);

	while(q--){
		scanf("%d%d%d",&k,&m,&r);
		vector<int>nodes;nodes.clear();
		for(int i=0;i<k;i++){
			scanf("%d",&u);
			nodes.push_back(u);
			sp[u]=1;
		}
		nodes.push_back(r);
		build_vtree(nodes);
		memset(dp,0,sizeof(dp));
		dp[0]=1;
		dfs2(r,0,0);
		ans=0;
		for(int i=0;i<=m;i++)ans=(ans+dp[i])%MOD;
		printf("%lld\n",ans);
		for(auto u:nodes)sp[u]=0;
		for(auto u:nodes)vt[u].clear();
	}
}