CF1770E Koxia and Tree

题目描述

给定一棵nnn个点的树,在kkk个位置上存在蝴蝶,我们需要给n−1n-1n−1条边定向,如果一条边的起点有蝴蝶且终点没有蝴蝶,那么蝴蝶将被移动到终点,我们会按照给定边的顺序移动,问最终所有蝴蝶的树上距离的和的期望,答案除于k(k−1)2\frac{k(k-1)}{2}2k(k−1)​,对998244353998244353998244353取模
k≤n≤300000k\le n\le 300000k≤n≤300000

题解

首先考虑第一个问题,我们怎么求kkk个点的距离和,可以随便钦定一个点为根,求出sizusiz_usizu​表示子树uuu中的指定结点个数,答案显然就是∑sizu(k−sizu)\sum siz_u(k-siz_u)∑sizu​(k−sizu​)
意识到每条边只会经历一遍,所以sizusiz_usizu​只会变化111,可以通过uuu到父亲的边的移动情况讨论得出
由于uuu到其父亲上的边未走过之前,uuu和其父亲有蝴蝶的概率是互不影响的,所以我们只需要知道uuu和其父亲上有蝴蝶的概率就可以算出Δans\Delta ansΔans,新的概率也可以动态维护pu=pfau=pu+pfau2p_u=p_{fa_u}=\frac{p_u+p_{fa_u}}{2}pu​=pfau​​=2pu​+pfau​​​

code\text{code}code

#include
#include
#include
#define ll long long
using namespace std;
const ll mod=998244353,inv2=499122177;
ll ksm(ll a,ll b)
{if(b==0) return 1;ll tmp=ksm(a,b>>1);if(b&1) return tmp*tmp%mod*a%mod;else return tmp*tmp%mod; 
}
void read(int &res)
{res=0;char ch=getchar();while(ch<'0'||ch>'9') ch=getchar();while('0'<=ch&&ch<='9') res=(res<<1)+(res<<3)+(ch^48),ch=getchar();
}
const int N=3e5+100;
int n,k,a[N+10],A[N+10],B[N+10];
vector e[N+10];
bool vis[N+10];
int fa[N+10],siz[N+10];
ll ans,p[N+10];
void dfs(int u)
{siz[u]=vis[u];for(int v:e[u]){if(v==fa[u]) continue;fa[v]=u,dfs(v),siz[u]+=siz[v];}ans+=1ll*siz[u]*(k-siz[u]);
}
int main()
{read(n),read(k);for(int i=1;i<=k;i++) read(a[i]),vis[a[i]]=true,p[a[i]]=1;for(int i=1;iif(fa[A[i]]==B[i]) swap(A[i],B[i]);ans+=p[A[i]]*(mod+1-p[B[i]])%mod*(1ll*(siz[B[i]]+1)*(k-siz[B[i]]-1)%mod+mod-1ll*siz[B[i]]*(k-siz[B[i]])%mod)%mod*inv2%mod,ans%=mod;ans+=p[B[i]]*(mod+1-p[A[i]])%mod*(1ll*(siz[B[i]]-1)*(k-siz[B[i]]+1)%mod+mod-1ll*siz[B[i]]*(k-siz[B[i]])%mod)%mod*inv2%mod,ans%=mod;p[A[i]]=p[B[i]]=(p[A[i]]+p[B[i]])*inv2%mod;}printf("%lld\n",ans*ksm(1ll*k*(k-1)%mod*inv2%mod,mod-2)%mod);return 0;
}