Luogu P5298 PKUWC2018 Minimax 题解 [ 紫 ] [ 树形 dp ] [ 线段树合并 ] [ 概率 dp ]

Minimax：线段树合并优化 dp 好题。

树形 dp

因为要求出每一个值的出现概率，首先我们可以想到一个很暴力的 dp 式子。

定义 \(dp_{i,j}\) 表示在节点 \(i\) 时，权值 \(j\) 的出现概率，设 \(l\) 表示左儿子，\(r\) 表示右儿子，则有如下转移：

• 当 \(j\) 在左儿子中时，\(dp_{i,j}\gets dp_{l,j}\times(p_i\times\sum_{k=1}^{j-1}dp_{r,k}+(1-p_i)\times\sum_{k=j+1}^{V}dp_{r,k})\)，理解的话就是对父亲节点选大的还是选小的进行分讨。
• 当 \(j\) 在右儿子中时，\(dp_{i,j}\gets dp_{r,j}\times(p_i\times\sum_{k=1}^{j-1}dp_{l,k}+(1-p_i)\times\sum_{k=j+1}^Vdp_{l,k})\)。

直接转移即可，时间复杂度 \(O(nV)\)。

线段树合并优化

显然原来的时间复杂度会炸掉，但是我们发现每个节点最开始时最多只有一个 dp 位置是有值的，所以我们考虑用这种均摊复杂度的线段树合并来优化这个 dp。

因为 dp 转移的时候需要用到前缀和后缀和，所以我们进行 merge 的时候记录节点 \(x,y\) 的前缀和 \(px,py\) 与后缀和 \(sx,sy\) 以及父亲节点的概率 \(p\)。

梳理一下 merge 的流程：

• 进入节点 \(x,y\)。
• 如果 \(x,y\) 其中之一是空树，则说明直接更新 dp 值即可。
  • 若 \(x\) 是空树，对应着上述 \(j\) 在右儿子中的转移方式，则我们对 \(y\) 的整颗子树内的 dp 值全部乘上 \((p\times\sum_{k=1}^{j-1}dp_{l,k}+(1-p)\times\sum_{k=j+1}^Vdp_{l,k})=(p\times px+(1-p)\times sx)\) 即可。这个可以用懒标记实现区间乘。
  • 若 \(y\) 是空树，对应着上述 \(j\) 在左儿子中的转移方式，则我们对 \(x\) 的整颗子树内的 dp 值全部乘上 \((p\times\sum_{k=1}^{j-1}dp_{r,k}+(1-p)\times\sum_{k=j+1}^Vdp_{r,k})=(p\times py+(1-p)\times sy)\) 即可。这个可以用懒标记实现区间乘。
• 否则就说明要递归合并，递归左右儿子的时候记得更新 \(sx,sy,px,py\) 的值。
• 最后将左右儿子的 dp 值加起来就是这个区间的 dp 值。

时间复杂度 \(O(n\log n)\)。

代码

#include <bits/stdc++.h>
#define fi first
#define se second
#define eb(x) emplace_back(x)
#define pb(x) push_back(x)
#define lc(x) (tr[x].ls)
#define rc(x) (tr[x].rs)
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ldb;
using pi=pair<int,int>;
const int N=300005;
const ll mod=998244353;
int n,fa[N],m=0,b[N],son[N][2],cd[N],p[N],ans[N];
ll qpow(ll a,ll b)
{
    ll res=1;
    while(b)
    {
        if(b&1)res=(res*a)%mod;
        b>>=1;
        a=(a*a)%mod;
    }
    return res;
}
int getrk(int x)
{
    return (lower_bound(b+1,b+m+1,x)-b);
}
struct Node{
    int ls,rs;
    ll dp,tag=1;
};
struct Segtree{
    Node tr[20*N];
    int root[N],tot=0;
    void pushup(int p)
    {
        tr[p].dp=(tr[lc(p)].dp+tr[rc(p)].dp)%mod;
    }
    void pushdown(int p)
    {
        if(tr[p].tag!=1)
        {
            tr[lc(p)].tag=(tr[lc(p)].tag*tr[p].tag)%mod;
            tr[rc(p)].tag=(tr[rc(p)].tag*tr[p].tag)%mod;
            tr[lc(p)].dp=(tr[lc(p)].dp*tr[p].tag)%mod;
            tr[rc(p)].dp=(tr[rc(p)].dp*tr[p].tag)%mod;
        }
        tr[p].tag=1;
    }
    void modify(int p,int v)
    {
        tr[p].dp=(tr[p].dp*1ll*v)%mod;
        tr[p].tag=(tr[p].tag*1ll*v)%mod;
    }
    void update(int &u,int ln,int rn,int x,ll k)
    {
        if(u==0)u=++tot;
        if(ln==rn){tr[u].dp+=k;return;}
        int mid=(ln+rn)>>1;
        if(x<=mid)update(lc(u),ln,mid,x,k);
        else update(rc(u),mid+1,rn,x,k);
        pushup(u);
    }
    int merge(int x,int y,int px,int py,int sx,int sy,int p)
    {
        if(x==0&&y==0)return 0;
        if(x==0)
        {
            modify(y,(1ll*p*px%mod+1ll*((1-p)%mod+mod)%mod*sx)%mod);
            return y;
        }
        if(y==0)
        {
            modify(x,(1ll*p*py%mod+1ll*((1-p)%mod+mod)%mod*sy)%mod);
            return x;
        }
        pushdown(x);pushdown(y);
        int lx=tr[lc(x)].dp,rx=tr[rc(x)].dp,ly=tr[lc(y)].dp,ry=tr[rc(y)].dp;
        tr[x].ls=merge(lc(x),lc(y),px,py,(sx+rx)%mod,(sy+ry)%mod,p);
        tr[x].rs=merge(rc(x),rc(y),(px+lx)%mod,(py+ly)%mod,sx,sy,p);
        pushup(x);
        return x;
    }
    void query(int u,int ln,int rn)
    {
        if(ln==rn){ans[ln]=tr[u].dp;return;}
        int mid=(ln+rn)>>1;
        pushdown(u);
        query(lc(u),ln,mid);
        query(rc(u),mid+1,rn);
    }
}tr1;
void dfs1(int u)
{
    if(son[u][0]==0)
    {
        tr1.update(tr1.root[u],1,m,getrk(p[u]),1);
        return;
    }
    if(son[u][1]==0)
    {
        dfs1(son[u][0]);
        tr1.root[u]=tr1.root[son[u][0]];
        return;
    }
    dfs1(son[u][0]);
    dfs1(son[u][1]);
    tr1.root[u]=tr1.merge(tr1.root[son[u][0]],tr1.root[son[u][1]],0,0,0,0,p[u]);
}
int main()
{
    //freopen("sample.in","r",stdin);
    //freopen("sample.out","w",stdout);
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin>>n;
    for(int i=1;i<=n;i++)cin>>fa[i];
    for(int i=1;i<=n;i++)
    {
        son[fa[i]][cd[fa[i]]]=i;
        cd[fa[i]]++;
    }
    for(int i=1;i<=n;i++)
    {
        cin>>p[i];
        if(cd[i])p[i]=p[i]*1ll*qpow(10000,mod-2)%mod;
        else b[++m]=p[i];
    }
    sort(b+1,b+m+1);
    m=unique(b+1,b+m+1)-b-1;
    dfs1(1);
    tr1.query(tr1.root[1],1,m);
    ll res=0;
    for(int i=1;i<=m;i++)res=(res+1ll*i*b[i]%mod*ans[i]%mod*ans[i]%mod)%mod;
    cout<<res;
    return 0;
}