第四届西安邮电大学acm-icpc校赛 热狗树

题目描述

 “我是番茄酱！”
“我是黄芥末酱！”
“合在一起就是——美式热狗上加的，那个！“
热狗树上的每个节点都涂有番茄酱或者黄芥末酱中的一种，这样热狗树就变得美味了~LiMn2O4构造了一颗热狗树，他想知道这棵树的美味程度。
一个热狗树的美味程度，定义为每个节点到其他和自己口味不一样的节点的最短距离之和的和，也就是任意两个口味不同的节点之间的路径长度和。请你求出这颗树的美味值，并且答案对998244353取模。

 

 

输入

第一行一个正整数N。
第二行有N个数，他们的值在[0,1]范围内取，其中1代表是涂了番茄酱，0代表是涂了黄芥末酱。
接下来N-1行，每一行有三个数a,b,w。代表节点a到b有边，路径的权值是w。输入数据保证是一棵树。
N≤100000,1≤a,b≤N,w<1 000 000 000

 

 

输出

输出答案对998244353取模后的结果。

 

样例输入

3
1 0 1
1 2 1
1 3 2

样例输出

8

 

提示

树（tree）是包含n（n>=0）个结点的有穷集，其中：
（1）每个元素称为结点（node）；
（2）有一个特定的结点被称为根结点或树根（root）。
（3）除根结点之外的其余数据元素被分为m（m≥0）个互不相交的集合T1，T2，……Tm-1，其中每一个集合Ti（1<=i<=m）本身也是一棵树，被称作原树的子树（subtree）。

 　　很裸的一个树形dp，sb的一直在用点算，其实直接以边的角度来看，每条边的贡献值就是，2ll*(子节点子树为1的数目*（总0数目-子节点子树为0的数目）*边权+子节点子树为0的数目*（总1数目-子节点子树为1的数目）*边权)

　　也就是看这条边两边为0和为1的情况。

#include<cstdio>
typedef long long ll;
const int N=100118,mod=998244353;
struct Side{
    int v,ne;
    ll w;
}S[N<<1];
int sn,head[N],size[N][2],a[N],num[2];
ll ans;
void init(int n)
{
    sn=0;
    ans=0;
    num[0]=num[1]=0;
    for(int i=0;i<=n;i++){
        head[i]=-1;
        size[i][0]=size[i][1]=0;
    }
}
void add(int u,int v,int w)
{
    S[sn].v=v;
    S[sn].w=1ll*w;
    S[sn].ne=head[u];
    head[u]=sn++;
}
void dfs(int u,int f)
{
    size[u][0]=(a[u]==0);
    size[u][1]=(a[u]==1);
    for(int i=head[u],v;~i;i=S[i].ne){
        v=S[i].v;
        if(v!=f){
            dfs(v,u);
            size[u][0]+=size[v][0];
            size[u][1]+=size[v][1];
            ans=(ans+2ll*size[v][0]*(num[1]-size[v][1])*S[i].w%mod)%mod;
            ans=(ans+2ll*size[v][1]*(num[0]-size[v][0])*S[i].w%mod)%mod;
        }
    }
}
int main()
{
    int n,u,v,w;
    while(~scanf("%d",&n)){
        init(n);
        for(int i=1;i<=n;i++){
            scanf("%d",&a[i]);
            num[0]+=(a[i]==0);
            num[1]+=(a[i]==1);
        }
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        dfs(1,0);
        printf("%lld\n",ans);
    }
    return 0;
}