CF1770E Koxia and Tree

题目描述

Imi has an undirected tree with $n$ vertices where edges are numbered from $1$ to $n-1$ . The $i$ -th edge connects vertices $u_i$ and $v_i$ . There are also $k$ butterflies on the tree. Initially, the $i$ -th butterfly is on vertex $a_i$ . All values of $a$ are pairwise distinct.

Koxia plays a game as follows:

• For $i=1,2,\dots ,n-1$ , Koxia set the direction of the $i$ -th edge as $u_i\to v_i$ or $v_i\to u_i$ with equal probability.
• For $i=1,2,\dots ,n-1$ , if a butterfly is on the initial vertex of $i$ -th edge and there is no butterfly on the terminal vertex, then this butterfly flies to the terminal vertex. Note that operations are sequentially in order of $1,2,\dots ,n-1$ instead of simultaneously.
• Koxia chooses two butterflies from the $k$ butterflies with equal probability from all possible $\frac{k(k-1)}{2}$ ways to select two butterflies, then she takes the distance ${}^{†}$ between the two chosen vertices as her score.

Now, Koxia wants you to find the expected value of her score, modulo $998244353^{‡}$ .

${}^{†}$ The distance between two vertices on a tree is the number of edges on the (unique) simple path between them.

${}^{‡}$ Formally, let $M=998244353$ . It can be shown that the answer can be expressed as an irreducible fraction $\frac{p}{q}$ , where $p$ and $q$ are integers and $q\not\equiv 0(\mathrm{mod}M)$ . Output the integer equal to $p\cdot q^{-1}\mathrm{mod}M$ . In other words, output such an integer $x$ that $0\le x<M$ and $x\cdot q\equiv p(\mathrm{mod}M)$ .

题意简述

一棵有 $n$ 个节点的树，其中有 $k$ 个节点上有一只蝴蝶。

• 你可以选择对每一条边随机的指定一个方向。
• 按照编号顺序依次选择每一条边，如果起点有蝴蝶，终点没有蝴蝶，则蝴蝶从起点飞到终点。
• 随机的选择两只蝴蝶，求他们距离的期望值。距离是两点之间的边数。

题解

如果不考虑蝴蝶的移动，这就是一个经典问题：求 $k$ 个蝴蝶两两距离的期望。

对于这个问题，可以枚举每条边，这条边把整棵树分成两份，若左边的点要走到右边的点去，则一定会经过这条边，对答案的贡献就是两边蝴蝶个数的乘积。最后再除上 $\frac{k(k-1)}{2}$ 就是最终的答案。

现在考虑蝴蝶的移动。再次枚举每条边，此时蝴蝶移动的可能性有三种：

• 蝴蝶在左边，且移动后的位置也在左边
• 蝴蝶在右边，且移动后的位置也在右边
• 蝴蝶在枚举的边的端点上

可以看到无论如何，左右两边的蝴蝶数量最多变化 $1$，且只有蝴蝶在枚举的边的端点上时才可能变化。

由于蝴蝶移动是概率性的，不妨设节点 $u$ 有蝴蝶的概率为 $p_u$，再设 $s_u$ 为子树 $u$ 中蝴蝶的数量。

按照编号由小到大枚举每条边 $(u,v)$，其中 $v$ 是 $u$ 的父亲，考虑求出这条边的贡献，分情况讨论：

• $u$ 和 $v$ 都有蝴蝶，此时的贡献为 $p_u\cdot p_v\cdot s_u\cdot (k-s_u)$
• $u$ 和 $v$ 都没有蝴蝶，此时的贡献为 $(1-p_u)\cdot (1-p_v)\cdot s_u\cdot (k-s_u)$
• $u$ 有蝴蝶，$v$ 没有蝴蝶，此时有 $\frac{1}{2}$ 的概率蝴蝶会从$u$ 飞到 $v$，贡献为 $p_u\cdot (1-p_v)\cdot \frac{s_u\cdot (k-s_u)+(s_u-1)\cdot (k-s_u+1)}{2}$
• $u$ 没有蝴蝶，$v$ 有蝴蝶，此时有 $\frac{1}{2}$ 的概率蝴蝶会从$v$ 飞到 $u$，贡献为 $(1-p_u)\cdot p_v\cdot \frac{s_u\cdot (k-s_u)+(s_u+1)\cdot (k-s_u-1)}{2}$

把以上贡献加起来就是边 $(u,v)$ 的贡献。

蝴蝶在边 $(u,v)$ 上有概率移动，$p_u$ 也会变化，考虑分四种情况讨论如何转移 $p_u$。

• $u$ 和 $v$ 都有蝴蝶，显然概率为 $p_u\cdot p_v$
• $u$ 和 $v$ 都没有蝴蝶，显然概率为 $0$
• $u$ 有蝴蝶，$v$ 没有蝴蝶，有 $\frac{1}{2}$ 的概率蝴蝶不会从$u$ 飞到 $v$，因此概率为 $\frac{p_u\cdot (1-p_v)}{2}$
• $u$ 没有蝴蝶，$v$ 有蝴蝶，有 $\frac{1}{2}$ 的概率蝴蝶会从$v$ 飞到 $u$，因此概率为 $\frac{(1-p_u)\cdot p_v}{2}$

总的概率为 $p_u\cdot p_v+\frac{p_u\cdot (1-p_v)}{2}+\frac{(1-p_u)\cdot p_v}{2}=\frac{p_u+p_v}{2}$

即转移时 $p_u\leftarrow \frac{p_u+p_v}{2}$

对于 $p_v$ 同理，也为 $\frac{p_u+p_v}{2}$，和 $p_u$ 相等。

别忘记最后答案要除以 $\frac{k(k-1)}{2}$。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod=998244353;
const int N=3e5+1;
int n,k;
int head[N],nxt[N<<1],to[N<<1],cnt,fa[N];
ll t[N],ans,p[N],s[N];
ll ksm(ll a,ll b)
{
    ll ans=1;
    while(b){
        if(b&1) ans=ans*a%mod;
        b>>=1;
        a=a*a%mod;
    }
    return ans;
}
void add(int u,int v)
{
    to[++cnt]=v;
    nxt[cnt]=head[u];
    head[u]=cnt;
}
void dfs(int u,int f)
{
    fa[u]=f;
    s[u]=p[u];
    for(int i=head[u];i;i=nxt[i]){
        if(to[i]!=f){
            dfs(to[i],u);
            s[u]+=s[to[i]];
        }
    }
}
int main()
{
    scanf("%d%d",&n,&k);
    for(int i=1,x;i<=k;i++) scanf("%d",&x),p[x]=1;
    for(int i=1,u,v;i<n;i++){
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dfs(1,0);
    ll inv2=ksm(2,mod-2);
    for(int i=1;i<=cnt;i+=2){
        int u=to[i],f=to[i+1];
        if(fa[u]!=f) swap(u,f);
        ans=((ans
        +p[u]*p[f]%mod*(k-s[u])%mod*s[u]%mod
        +(1-p[u])*(1-p[f])%mod*(k-s[u])%mod*s[u]%mod
        +p[u]*(1-p[f])%mod*((k-s[u])*s[u]%mod+(k-s[u]+1)*(s[u]-1)%mod)%mod*inv2%mod
        +(1-p[u])*p[f]%mod*((k-s[u])*s[u]%mod+(k-s[u]-1)*(s[u]+1)%mod)%mod*inv2%mod
        )%mod+mod)%mod;
        p[u]=p[f]=(p[u]+p[f])*inv2%mod;
    }
    cout<<ans*ksm(1ll*k*(k-1)/2%mod,mod-2)%mod;
}