西电校赛2020D题——双倍快乐

题意：

给你一个树，树上每个人有一个价值。定义一个快乐值：

1. 叶子结点的快乐值为价值
2. 其他结点的快乐值为儿子快乐值值和加上自己的价值

树上每条边要么是黑的要么是白的，如果一个结点既有黑边又有白边，那么它的价值就会翻倍。

现在允许你修改任意一条边的颜色，也可以不修改，问你能达到的最大全体快乐的总和是多少。

这题本来想 DP 做的，在队友的提示下可以从考虑修改每一条边造成的贡献入手。

每个点的快乐值实际上就是子树权值和，因此每个点对答案的贡献就是它的价值乘上深度。

如果修改一条边的颜色，那么这条边的两个端点有如下三种可能：

1. 权值翻倍
2. 无事发生
3. 权值减半

且两个端点的贡献可以独立考虑。

要判断是哪种情况，只需统计每条边原先连接的黑白边个数。假设从 A 边修改成了 B 边，那么：

1. 如果原先没有 B 边且 A 边至少有两个，那么价值翻倍
2. 如果原先只有 1 个 B 边且有 A 边，那么价值减半
3. 否则无事发生

可以先统计原先的快乐值总和，然后维护最大的可以得到的快乐值增加量。这个增加量初值为 0 ，表示可以不改。最后输出原先总和 + 最大增量即可。

代码：

#include<bits/stdc++.h>
using namespace std;
struct edg{
        int nxt,y,v;
}e[210000];
int lk[110000],lt;
int et[110000][2];
void ist(int x,int y,int z){
        e[++lt]=(edg){lk[x],y,z};
        lk[x]=lt;
        e[++lt]=(edg){lk[y],x,z};
        lk[y]=lt;
        et[x][z]++;
        et[y][z]++;
}
int n;
long long a[110000];
int dp[110000];
long long bow;
long long ans;
long long ccl(int x,int y){
        y^=1;
        if(et[x][y]==0 && et[x][y^1]>1){
                return a[x]*dp[x];
        }
        if(et[x][y]==1 && et[x][y^1]>0){
                return -a[x]*dp[x];
        }
        return 0;
}
void dfs(int x,int y,int z){
        dp[x]=z;
        for(int i=lk[x];i;i=e[i].nxt){
                if(e[i].y==y){
                        continue;
                }
                dfs(e[i].y,x,z+1);
                ans=max(ans,ccl(x,e[i].v)+ccl(e[i].y,e[i].v));
        }
        if(et[x][0] && et[x][1]){
                bow+=a[x]*2*dp[x];
        }
        else{
                bow+=a[x]*dp[x];
        }
}
int main(){
        scanf("%d",&n);
        for(int i=1;i<=n;i++){
                scanf("%lld",&a[i]);
        }
        int l,r,v;
        for(int i=1;i<n;i++){
                scanf("%d%d%d",&l,&r,&v);
                ist(l,r,v);
        }
        dfs(1,0,1);
        printf("%lld\n",bow+ans);
        return 0;
}