题解：P16365 [OOI 2026] Tasks from Sasha

· · 题解

拜谢大龙 @wzc6666 给我讲这题。

首先稍加思考可以得到一个 O(nk) 的暴力 dp。设 f_{u,i} 表示 u 子树内的点来填 1 到 fa_u 上的任务点的方案数，转移分讨 u 是否是关键点。

如果 u 不是关键点，u 自己有 i+1 种颜色可选，再乘所有儿子的贡献。

如果 u 是关键点，那么分讨 u 是否被染成自己的颜色。

如果染上，那么他的所有儿子都可以给 f_{v,i+1} 的贡献。

如果 u 没有被染成自己的颜色，那么我们就需要从她的至少两个子树中染出 u 的颜色，我们用所有情况减去选 0 个或 1 个子树的情况。0 个子树也就是说 u 有 i+1 种选择也就是上边的所有可选颜色以及 0。1 个子树我们枚举一个子树，钦定她出现过 u 的颜色，贡献为 f_{v,i+1}-f_{v,i}，其他儿子子树均不能出现 u 的颜色。

暴力转移式子为：

f_{u,i}= \begin{cases} (i+1)\displaystyle\prod_{v\in son_u}f_{v,i},&u\notin important\\ (i+2)\displaystyle\prod_{v\in son_u}f_{v,i+1}-(i+1)\displaystyle\prod_{v\in son_u}f_{v,i}-(i+1)\displaystyle\sum_{v\in son_u}\left(f_{v,i+1}-f_{v,i}\right)\prod_{\substack{w\in son_u\\w\ne v}}f_{w,i},&u\in important \end{cases}

然后我们注意到 k 很小，想要建出来类似虚树的东西，但是发现建完没法转移。我们先把任务点称作红点，非任务点称为白点。我们把每个白点连通块缩成一个点，这样点数降到了 O(k) 吗？如果说一个菊花图，根节点是红点其它都是白点就完蛋了，所以我们一会单独考虑叶子，先考虑缩完的中间非叶子节点怎么转移。

为了方便书写，我们设 g_{u,i} = \prod_{v\in son_u} f_{v,i}。

白色中间非叶子节点，每个点都有 (i+1) 种选择，转移即为 f_{u,i} = (i+1)^{sz} g_{u,i}，其中 sz 表示这个白点实际代表的大小。

红色中间非叶子节点没区别。

接下来考虑叶子，白叶子按照上边的转移为 f_{u,i} = (i+1)^{sz}，红叶子 f_{u,i} = 1。

此时我们发现啥也没优化，时间复杂度甚至多个快速幂的 \log。

此时我们发现，所有白叶子本质上只有大小的区别，所以我们记 cnt_i 表示原来大小为 i 的白叶子有几个，所以一共的白叶子大小为 S=\sum_i i cnt_i。刚才式子变成 g_{u,i} = (i+1)^S \prod_{v\in R_u} f_{v,i}。

现在底下紧挨着白叶子的红色中间节点的只选一个子树的我们想快速算，f_{u,i}=(i+2)g_{u,i+1}-(i+1)g_{u,i}-(i+1)E，E 就是我们想求的恰好一个儿子使用 u 的新颜色。

现在考虑一个大小为 s 的白叶子，不选的贡献是 (i+1)^s，选的贡献是 (i+2)^s-(i+1)^s，但是 g_{u,i} 里边算过了她的贡献，所以除掉原贡献，乘上新贡献，相同大小的白叶子贡献相同，因此可以合并计算，再乘上该大小的白叶子数量。直接乘个数 E_1 = g_{u,i} \sum_s cnt_s \frac{(i+2)^s-(i+1)^s}{(i+1)^s}。

其他儿子我们正常钦定即可，即 E_2 =\sum_{v\in R_u}(f_{v,i+1}-f_{v,i})\prod_{\substack{w\in son_u\\w\ne v}}f_{w,i}。

最后的转移为 f_{u,i}=(i+2)g_{u,i+1}-(i+1)g_{u,i}-(i+1)(E_1+E_2)。

然后发现这个不同叶子大小种类是等差数列量级 O(\sqrt n)。整个树需要的非叶子个数只有 O(k) 个，所以时间复杂度 O(k \sqrt{kn} \log mod)。幂次和逆元直接算肯定过不去。但是我们非叶子的都可以暴力算，叶子的幂底数 O(k)，指数 O(n)，所以我们可以预处理光速幂计算即可。

设第 u 个非叶子节点挂着的白叶子总大小为 S_u，它的不同大小种数为 O(\sqrt{S_u})。

因为非叶子有 O(k) 个，且所有白叶子的总大小不超过 n，根据柯西不等式：\sum_u\sqrt{S_u}\le \sqrt{O(k)\sum_u S_u}=O(\sqrt{kn})。

再枚举 O(k) 个 DP 状态，得到：O(k\sqrt{kn})。

最终时间复杂度 O(n + k\sqrt n + k^2 \log mod + k \sqrt{nk})。

Code.

#include<bits/stdc++.h>
const int mod=998244353;
using namespace std;
#define int long long
#define main(x) MiAn();int mIaN();main(){ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);int t=1;if(strlen(#x))cin>>t;while(t--)mIaN();cerr<<"Time: "<<1.0*clock()/CLOCKS_PER_SEC<<"s\n";}int mIaN()
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define per(i,a,b) for(int i=(a);i>=(b);i--)
#define mem(a,b) memset((a),(b),sizeof(a))
#define all(x) (x).begin(),(x).end()
#define abs_(x) ((x)<0?-(x):(x))
#define sz(x) ((int)(x).size())
#define ull unsigned long long
#define pb push_back
#define fi first
#define se second
const int block=911;
int pw1[2005][block];
int pw2[2005][1100];
int ipw1[2005][block];
int ipw2[2005][1100];
int qpow(int a,int b){
    int res=1;
    while(b){
        if(b&1) res=res*a%mod;
        a=a*a%mod;
        b>>=1;
    }
    return res;
}
int fastpow(int a,int b){return pw1[a][b%block]*pw2[a][b/block]%mod;}
int fastinv(int a,int b){return ipw1[a][b%block]*ipw2[a][b/block]%mod;}
int n,k;
int tot;
int B[1000005];
int fa[1000005];
int sc[1000005];
int sz[1000005];
int sum[1000005];
bool col[1000005];
bool inp[1000005];
int f[4005][2005];
int dp[4005][2005];
int id[1000005];
int cnt[1000005];
vector<int>mp[1000005];
int faid_to_sc[1000005];
vector<int>newedge[1000005];
void initpow(int n){
    rep(a,1,k+2){
        int ia=qpow(a,mod-2);
        pw1[a][0]=pw2[a][0]=1;
        ipw1[a][0]=ipw2[a][0]=1;
        rep(j,1,block-1){
            pw1[a][j]=pw1[a][j-1]*a%mod;
            ipw1[a][j]=ipw1[a][j-1]*ia%mod;
        }
        int x=pw1[a][block-1]*a%mod;
        int y=ipw1[a][block-1]*ia%mod;
        rep(j,1,n/block){
            pw2[a][j]=pw2[a][j-1]*x%mod;
            ipw2[a][j]=ipw2[a][j-1]*y%mod;
        }
    }
}
void dfs(int u,int faa){
    sum[u]=sum[faa]+col[u];
    id[u]=++tot;
    for(int v:mp[u]){
        if(v==faa) continue;
        if(!col[v]&&!(sz(mp[v])^1)) continue;
        dfs(v,u);
    }
    vector<int>types;
    int qwq=0;
    for(int v:mp[u]){
        if(v==faa) continue;
        if(!col[v]&&!(sz(mp[v])^1)){
            if(!cnt[sz[v]]) types.pb(sz[v]);
            cnt[sz[v]]++;
            qwq+=sz[v];
        }
    }
    rep(i,0,sum[u]){
        dp[id[u]][i]=fastpow(i+1,qwq);
        for(int v:mp[u]){
            if(v==faa) continue;
            if(!col[v]&&!(sz(mp[v])^1)) continue;
            dp[id[u]][i]=dp[id[u]][i]*f[id[v]][i]%mod;
        }
    }
    rep(i,0,sum[u]-col[u]){
        if(!col[u]) f[id[u]][i]=dp[id[u]][i]*fastpow(i+1,sz[u])%mod;
        else{
            int ans=((i+2)*dp[id[u]][i+1]-(i+1)*dp[id[u]][i]%mod+mod)%mod;
            int zero=0,dyy=1,res=0,awa=0;
            for(int v:mp[u]){
                if(v==faa) continue;
                if(!col[v]&&!(sz(mp[v])^1)) continue;
                int a=f[id[v]][i];
                int b=f[id[v]][i+1];
                if(!a) zero++,awa=b;
                else dyy=dyy*a%mod,res=(res+(b-a+mod)*qpow(a,mod-2))%mod;
            }
            for(int s:types) res=(res+cnt[s]*(fastpow(i+2,s)*fastinv(i+1,s)%mod-1+mod))%mod;
            if(!zero) res=res*dp[id[u]][i]%mod;
            else if(!(zero^1)) res=awa*dyy%mod*fastpow(i+1,qwq)%mod;
            else res=0;
            f[id[u]][i]=(ans-(i+1)*res%mod+mod)%mod;
        }
    }
    for(int s:types) cnt[s]=0;
}
int main(){
    cin>>n>>k;
    initpow(n);
    rep(i,1,k){int x;cin>>x;inp[x]=1;}
    rep(i,2,n){cin>>fa[i];mp[i].pb(fa[i]);mp[fa[i]].pb(i);}
    rep(i,1,n) sc[i]=((i^1)&&!(inp[i]|inp[fa[i]])?sc[fa[i]]:++fa[0]),sz[sc[i]]+=!inp[i];
    rep(i,1,n) if(inp[i]) col[sc[i]]=1;
    rep(u,1,n){for(int v:mp[u]){if(sc[u]^sc[v]) newedge[sc[u]].pb(sc[v]);}}
    n=fa[0];
    rep(u,1,n){sort(all(newedge[u]));newedge[u].erase(unique(all(newedge[u])),newedge[u].end());mp[u].clear();for(int v:newedge[u]){mp[u].pb(v);}}
    dfs(1,0);
    cout<<f[1][0];
    return 0;
}