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();
    }
}