【NOIP2017提高组模拟6.30】tty's home

Description

Input

Output

一行，答案mod998244353后的值

Sample Input

input 1:
5
1 1 1 1 1
1 2
2 3
3 4
4 5
input 2:
5
0 1 0 1 0
1 2
2 3
3 4
4 5

Sample Output

output 1:
15
output 2:
12

Solution

这道水题居然花了我一个小时

答案用总方案数减去没有选最喜欢的那些的方案数
设f[x]表示一定选点x的方案数
设y为x的儿子
f[x]+=f[y]+f[x]∗f[y]
此时，f[x]中有的是在子树y之前的所有子树的方案，f[x]*f[y]就是之前的每一种与y子树里的匹配，+f[y]就是x点本身与y子树中的所有方案的匹配
每次递归的时候都要统计答案，因为f[x]一定选了x这个点
减的时候方法相同，如果一个点是被喜欢的点，就只递归进去算它儿子们的答案而不统计自身的答案，自身的答案也不与父亲和兄弟们匹配

Code

#include<cstdio>#include<algorithm>#include<cstring>#define fo(i,a,b) for(int i=a;i<=b;i++)#define N 101000#define cl(a) memset(a,0,sizeof(a))#define mo 998244353#define ll long longusing namespace std;int a[N],mx=-2147483647,last[N],next[N*10],to[N*10],tot=0,bz[N],n;ll f[N],ans;void putin(int x,int y){    next[++tot]=last[x];last[x]=tot;to[tot]=y;}void dg(int x){    bz[x]=1;    for(int i=last[x];i;i=next[i])    {        int y=to[i];        if(bz[y]) continue;        dg(y);        f[x]=(f[x]+f[y]+(f[x]*f[y])%mo)%mo;    }    f[x]++;ans+=f[x];}void dg1(int x){    bz[x]=1;    for(int i=last[x];i;i=next[i])    {        int y=to[i];        if(bz[y]) continue;        dg1(y);        if(a[y]!=mx) f[x]=(f[x]+f[y]+(f[x]*f[y])%mo)%mo;    }    if(a[x]!=mx) f[x]++,ans=(ans-f[x]+mo)%mo;}int main(){    scanf("%d",&n);    fo(i,1,n) scanf("%d",&a[i]),mx=max(mx,a[i]);    fo(i,1,n-1)    {        int x,y;scanf("%d%d",&x,&y);        putin(x,y);putin(y,x);    }    dg(1);    cl(bz);cl(f);    dg1(1);    printf("%d",ans);}