Tree CodeForces - 1111E (树,计数,换根)

大意: 给定树, 多组询问, 每个询问给出一个点集$S$, 给定$m, r$, 求根为$r$时, $S$的划分数, 满足

• 每个划分大小不超过$m$
• 每个划分内不存在一个点是另一个点的祖先

设点$x$的祖先包括x的个数为$f[x]$, 按$f$排序后, 有转移

$$dp[i][j]=dp[i-1][j-1]+(j-f[i]+1)dp[i-1][j]$$

dp[i][j]为前i个, 划分为j组的方案数

所以讨论一下$r$与$1$的位置关系求出f就行了

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <math.h>
#include <set>
#include <map>
#include <queue>
#include <string>
#include <string.h>
#define REP(i,a,n) for(int i=a;i<=n;++i)
#define PER(i,a,n) for(int i=n;i>=a;--i)
#define hr putchar(10)
#define pb push_back
#define lc (o<<1)
#define rc (lc|1)
#define mid ((l+r)>>1)
#define ls lc,l,mid
#define rs rc,mid+1,r
#define x first
#define y second
#define io std::ios::sync_with_stdio(false)
#define endl '\n'
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int P = 1e9+7, INF = 0x3f3f3f3f;
ll gcd(ll a,ll b) {return b?gcd(b,a%b):a;}
ll qpow(ll a,ll n) {ll r=1%P;for (a%=P;n;a=a*a%P,n>>=1)if(n&1)r=r*a%P;return r;}
ll inv(ll x){return x<=1?1:inv(P%x)*(P-P/x)%P;}
//head

const int N = 1e5+10;
int n, q;
vector<int> g[N];
int top[N], fa[N], dep[N], son[N];
int L[N], R[N], sz[N];

int calc(int x, int y, int f) {
	int pre = 0, lca;
	while (top[x]!=top[y]) {
		if (dep[top[x]]<dep[top[y]]) swap(x,y);
		pre = top[x];
		x = fa[pre];
	}
	if (x==y) lca=x;
	else lca=dep[x]<dep[y]?x:y;
	if (lca!=f) return 0;
	if (fa[pre]==f) return pre;
	return son[f];
}

void dfs(int x, int f, int d) {
	L[x]=++*L,sz[x]=1,fa[x]=f,dep[x]=d;
	int mx = -1;
	for (int y:g[x]) if (y!=f) {
		dfs(y,x,d+1),sz[x]+=sz[y];
		if (mx<sz[y]) mx=sz[y],son[x]=y;
	}
	R[x]=*L;
}
void dfs2(int x, int tf) {
	top[x]=tf;
	if (son[x]) dfs2(son[x],tf);
	for (int y:g[x]) if (y!=fa[x]&&y!=son[x]) dfs2(y,y);
}
int c[N];
void add(int x, int v) {
	for (; x<=n; x+=x&-x) c[x]+=v;
}
void update(int l, int r, int v) {
	add(l,v),add(r+1,-v);
}
int query(int x) {
	int r = 0;
	for (; x; x^=x&-x) r+=c[x];
	return r;
}
struct _ {
	int x,c,f;
} a[N];
int dp[N][322];

int main() {
	scanf("%d%d", &n, &q);
	REP(i,2,n) {
		int u, v;
		scanf("%d%d", &u, &v);
		g[u].pb(v),g[v].pb(u);
	}
	dfs(1,0,0),dfs2(1,1);
	while (q--) {
		int k, m, r;
		scanf("%d%d%d", &k, &m, &r);
		REP(i,1,k) { 
			scanf("%d", &a[i].x);
			if (a[i].x==r) {update(1,n,1);continue;}
			a[i].c = calc(a[i].x,r,a[i].x);
			if (!a[i].c) update(L[a[i].x],R[a[i].x],1);
			else update(1,n,1),update(L[a[i].c],R[a[i].c],-1);
		}
		REP(i,1,k) a[i].f = query(L[a[i].x]);
		sort(a+1,a+1+k,[](_ a,_ b){return a.f<b.f;});
		if (a[k].f<=m) { 
			dp[0][0] = 1;
			REP(i,1,k) REP(j,a[i].f,m) {
				dp[i][j] = (dp[i-1][j-1]+(ll)(j-a[i].f+1)*dp[i-1][j])%P;
			}
			ll ans = 0;
			REP(i,a[k].f,m) ans+=dp[k][i];
			printf("%d\n", int(ans%P));
			REP(i,1,k) REP(j,a[i].f,m) dp[i][j]=0;
		}
		else puts("0");
		REP(i,1,k) {
			if (a[i].x==r) update(1,n,-1);
			else if (!a[i].c) update(L[a[i].x],R[a[i].x],-1);
			else update(1,n,-1),update(L[a[i].c],R[a[i].c],1);
		}
	}
}