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

Description

Input

Output

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

Data Constraint

Solution

• 正难则反，这题我们可以转换一下思路。姑且称 可能 符合条件的一种方案为 一个匹配 。

• 设 A= 不加限制的匹配数（总量）， B= 不符合条件的匹配数（部分），则答案即为 A-B 。

• 那么我们如何求出 A 和 B 呢？可以用到 树形DP 的方法。

• 设 $F_i$ 表示以 $i$ 为根的子树所产生的匹配数。则有转移方程式：

  $$F_i=F_i*F_j+F_j（j∈Son_i）$$ （逐一组合）

• 处理完上述式子后，再有：$F_i=F_i+1 \ $ （单独自己也是一种匹配）

• 最后可以求出 A ：

  $$A=\sum_{i=1}^{n}F_i$$

• 求 B 也同理，不过当 $j$ 为最值点时不转移 $i$ 、当 $i$ 为最值点时 $F_i$ 不 +1 。

• 最后 A、B 相减即可得出答案，总时间复杂度是 $O(N)$ 。

Code

#include<cstdio>
#include<cstring>
using namespace std;
const int N=100001,mo=998244353;
int tot,mx=-2147483647;
int first[N],next[N<<1],en[N<<1];
int a[N];
long long f[N],g[N];
long long ans;
inline int read()
{
    int X=0,w=1; char ch=0;
    while(ch<'0' || ch>'9') {if(ch=='-') w=-1;ch=getchar();}
    while(ch>='0' && ch<='9') X=(X<<3)+(X<<1)+ch-'0',ch=getchar();
    return X*w;
}
inline void insert(int x,int y)
{
    next[++tot]=first[x];
    first[x]=tot;
    en[tot]=y;
}
inline void dfs1(int x,int y)
{
    for(int i=first[x];i;i=next[i])
        if(en[i]!=y)
        {
            dfs1(en[i],x);
            (f[x]+=f[x]*f[en[i]]%mo+f[en[i]])%=mo;
        }
    ans+=++f[x];
}
inline void dfs2(int x,int y)
{
    for(int i=first[x];i;i=next[i])
        if(en[i]!=y)
        {
            dfs2(en[i],x);
            if(a[en[i]]<mx) (g[x]+=g[x]*g[en[i]]%mo+g[en[i]])%=mo;
        }
    if(a[x]<mx) (ans+=mo-++g[x])%=mo;
}
int main()
{
    int n=read();
    for(int i=1;i<=n;i++)
        if((a[i]=read())>mx) mx=a[i];
    for(int i=1;i<n;i++)
    {
        int x=read(),y=read();
        insert(x,y);
        insert(y,x);
    }
    dfs1(1,0),dfs2(1,0);
    printf("%lld",ans);
    return 0;
}