Description

EZ同学家里非常富有，但又极其的谦虚，说话又好听，是个不可多得的人才。
EZ常常在假期环游世界，他准备去N（N<=100000）个国家之多，一些国家有航线连接，由于EZ同学有一定的强迫症，任意两个国家之间都能通过航路直接或间接到达，并且这样的路径仅有一种。（简单来说，这些国家构成了一棵树）
由于EZ是C国人，因此将C国（1号国家）作为整棵树的根
每个国家有一个旅游热度A[i]和影响力D[i]。由于目的地有点多，为了避免选择困难症，他给每个国家设置了一个向往值F[i]，它等于所有的A[j]之和，满足i国在j国向C国走D[j]步的路径上（经过一条航路算一步，i=j也会被统计，如果D[j]步超过了C国，则超出部分不用管）。
LYD同学家里有矿，富有程度与EZ不相上下，但他却在宅与现充间摇摆不定。某次机缘巧合，EZ外出旅游刺激了LYD，他决定也要开始旅游。为了避免又被判高重复率导致被取消资格，他将EZ的旅游地图略微做了一点调整，每条航路将有一定的概率出现。
现在他有Q个询问，每次询问某个国家所在的联通块（由于每条边是一定概率出现，因此它所在的联通块可以是很多种）中所有国家的F[i]值的和的平方的期望（对998244353取模），以此来决定他旅游的目的地。但他极其厌恶繁琐的计算，于是他找到了能算出圆周率并将它倒背下来的你，答应给你丰厚的报酬。家里没矿，老爸也不是X达集团老总的你决定接受他的任务。

Input

第一行1个正整数N，表示国家数。
接下来N行，第i行两个非负整数A[i],D[i]，表示国家i的旅游热度以及影响力
接下来N-1行，每行三个非负整数x,y,v，x,y为这条航路连接的两个国家，v为这条航路出现的概率。（注意每个在EZ的地图中是没有出现概率的说法的，因此每个国家的F值与边的出现概率无关）
接下来一行一个正整数Q，表示询问数
接下来Q行，每行一个正整数x，表示询问国家x所在的联通块中所有国家的F[i]值的和的平方的期望（对998244353取模）。

Output

Q行，每行一个非负整数表示这次询问的答案。

Sample Input

Sample Input1:
5
2 1
1 0
3 1
3 2
1 4
1 2 1
1 3 1
2 4 1
2 5 0
3
1
3
5

Sample Input2:
4
5 1
3 0
10 2
4 1
1 2 74017368
1 3 59531864
2 4 25036401
3
4
2
3

Sample Output

Sample Output1:
400
400
1

样例解释： 
可算出各国的F值分别为9,5,3,3,1
1,2,3,4必定在同一个联通块中，5在另一个联通块
1,2,3,4所在联通块F值和为20，因此答案是400
5所在的联通块F值和为1，因此答案是1

Sample Output2:
988451137
606447316
733454972
样例解释： 
可算出各国的F值分别为15,7,10,4
根据期望的定义，可计算出答案

Data Constraint

对于100%的数据，1<=N,Q<=200000，0<=A[i],v<998244353,0<=D[i]<N

Hint

题解

刚看到题目的时候一脸懵逼，那按流程来
题目大意：给出一棵树，每个节点有两个值 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)²+(1-p)a²=pa²+pb²+pab+a²-pa²=p(b+ab)+a²
这样对于每个点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 }

[树上差分][子树求和][树形dp] Jzoj P5911 Travel