[树上差分][子树求和][树形dp] Jzoj P5911 Travel
题解
- 刚看到题目的时候一脸懵逼,那按流程来
- 题目大意:给出一棵树,每个节点有两个值 Ai,Di,表示这个节点能给它的 0~Di 级祖先的 F 值 贡献 Ai。然后给每条边一个出现概率,每次询问某个节点的联通块的 F 值的和的 平方的期望
- 首先考虑怎么求f[i],由于是在i向上跳D[i]级都要加上这个值,当然暴力太慢,考虑一下树上差分
- 在i点打上+1标记,然后再D[i]+1的祖先上打-1标记,就是子树求和可以解决的了
- 考虑如何计算和的平方,这个完全平方公式显然就是两两相乘再相加嘛,那么就可以每次操作从询问的国家开始dfs
- 用树形dp来求解,如果现在又两个要合并的联通块,值分别为a、b,然后成功的概率是p,答案:p(a+b)2+(1-p)a2=pa2+pb2+pab+a2-pa2=p(b+ab)+a2
- 这样对于每个点i,维护以它为根的子树中期望和的平方G[i],以及期望和H[i],按照上面从儿子合并到父亲即可
- 一直合并到根就是当前询问点的答案,这样一次询问的复杂度是O(N)的
- 那么这样的时间复杂度就是O(NQ),只有30分
- 考虑一下不用每次询问都dfs的做法,我们不妨一直以1为根
- 那么对于一个询问i,只用知道i的子树的G和H,和i以外部分的G和H
- 这样的话就可以从父亲的G和H和其他兄弟的G和H转移过来,就可以了,就类似与一次换根操作
代码
1 #include <cstdio> 2 #include <iostream> 3 #include <cstring> 4 const int N=2e5+10,M=4e5+10,mo=998244353; 5 using namespace std; 6 struct edge {int to,from,v;}e[M]; 7 int a[N],d[N],head[N],n,Q,x,cnt,p[N],tot; 8 long long f[N],G[N],H[N],g[N],h[N]; 9 void insert(int x,int y,int v) { e[++cnt].to=y; e[cnt].from=head[x]; e[cnt].v=v; head[x]=cnt; } 10 void dfs1(int x,int fa) 11 { 12 int k=p[max(tot-d[x],0)]; 13 p[++tot]=x,f[x]+=a[x],f[k]-=a[x]; 14 for (int i=head[x];i;i=e[i].from) 15 if (e[i].to!=fa) dfs1(e[i].to,x),f[x]+=f[e[i].to]; 16 f[x]=(f[x]%mo+mo)%mo,p[tot--]=0; 17 } 18 void dfs2(int x,int fa) 19 { 20 g[x]=f[x],h[x]=f[x]*f[x]%mo; 21 for (int i=head[x];i;i=e[i].from) 22 if (e[i].to!=fa) 23 dfs2(e[i].to,x), 24 h[x]=(h[x]+(2*g[x]*g[e[i].to]+h[e[i].to])%mo*e[i].v%mo)%mo, 25 g[x]=(g[x]+g[e[i].to]*e[i].v)%mo; 26 } 27 void dp(int x,int fa) 28 { 29 for (int i=head[x];i;i=e[i].from) 30 if (e[i].to!=fa) 31 { 32 long long a=(G[x]-g[e[i].to]*e[i].v%mo+mo)%mo,b=(H[x]-(a*g[e[i].to]%mo*2%mo+h[e[i].to])%mo*e[i].v%mo+mo)%mo; 33 H[e[i].to]=(h[e[i].to]+(g[e[i].to]*a%mo*2+b)%mo*e[i].v%mo)%mo; 34 G[e[i].to]=(g[e[i].to]+a*e[i].v)%mo; 35 dp(e[i].to,x); 36 } 37 } 38 int main() 39 { 40 freopen("travel.in","r",stdin),freopen("travel.out","w",stdout); 41 scanf("%d",&n); 42 for (int i=1;i<=n;i++) scanf("%d%d",&a[i],&d[i]); 43 for (int i=1,x,y,z;i<n;i++) scanf("%d%d%d",&x,&y,&z),insert(x,y,z),insert(y,x,z); 44 dfs1(1,0),dfs2(1,0),G[1]=g[1],H[1]=h[1],dp(1,0); 45 scanf("%d",&Q); 46 while (Q--) scanf("%d",&x),printf("%lld\n",H[x]); 47 }