树形DP+树上路径贡献

题目

一棵树有n个节点，每个节点都同有一个符号+或者-，表示该节点是正极节点还是负极节点。如果从正极走到负极，或者从负极走到正极，都会受到1点伤害。
现在需要知道走过所有路径后会受到的总伤害是多少？树上任意2点存在唯一的路径。需要计算所有任意2点的路径的伤害和。

注意：从u到，和从v到u视为同一条路径。

输入描述

第一行输入一个整数n（1≤n ≤ 10^5 )
第二行输入一个长度为1的仅由+和-组成的字符串。
接下来n-1 行，每行输入两个整数 u,v（1≤u,v≤n）表示节点
u和节点v之间有一条边。

输出描述

输出一个整数表示受到的总伤害。

样例1

输入

3
++-
1 2
1 3

输出
2

说明
路径1-2 （2-1）不会受到伤害。
脂径1-3（3-1）会受到1点伤害。
路径2-1-3（3-1-2）会受到1点害。
因此答案为2。

题目分析

典型的树形DP，换根可以解决路径去重的问题。现在考虑怎么计算贡献。
我们可以很容易想到一个f[i]: i的子树上所有经过点i的路径的伤害值的和。那么 $\sum\limits_{1}^{n} f[i]$ 就是答案

在DFS时到1的时候，考虑1-3的边，在考虑2和3的子树合并的时候计算f[1], 当前增加的路径是，根(1,2,4,5)到3点所有子树(3, 6)的路径。因为每条路径一定会经过1。对每条路径进行拆分， $(X - Y) = (X - 1) + (1 - Y)$
X是(1,2,4,5)，Y是(3, 6)。
每个X会和每个Y匹配出一条路径，那么对于X-1会被计算Y次。对于每个根u，我们求出它的字树到根u的路径和 $f1[u] = \sum\limits^{x=所有的子节点} (x 到u的路径伤害)$
那么这个公式
$(X - Y) = (X - 1) + (1 - Y)$
转化成
$f[u] = f1[u] * son[to] + f1[to] * son[to]$

f1[to]在计算前，应该处理u-to这条边的伤害。

代码如下：

#include <iostream>

using namespace std;

char ch[100005];
bool vis[100005];
vector<int> G[100005];
long long f[100005], f1[100005], son[100005];

void Dfs(int u, int fa){
    son[u] = 1;
    for (auto to : G[u]) {
        if (to != fa) {
            Dfs(to, u);
            // 处理边u-to的伤害
            if (ch[u] != ch[to]) {
                f1[to] += son[to];
            }
            // 计算贡献
            f[u] += f1[u] * son[to] + f1[to] * son[u];
            // 换根
            f1[u] += f1[to];
            son[u] += son[to];
        }
    }
}

int main() {
    
    int  n; scanf("%d", &n);
    scanf("%s", ch+1);
    for(int i=1; i<n; i++) {
        int a, b; scanf("%d%d", &a, &b);
        G[a].push_back(b);
        G[b].push_back(a);
    }
    long long ans = 0;
    Dfs(1, 0);
    for (int i=1; i<=n; i++) {
        ans += f[i];
    }
    printf("%lld\n", ans);
    
    return 0;
}
/*
7
++--+-+
1 2
2 3
2 4
3 5
1 6
6 7
 
5
+-++-
1 2
2 3
2 4
1 5
 
3
+--
 1 2
 1 3
 
 8
+-++-+-+
 1 2
 2 3
 3 4
 3 5
 1 6
 6 7
 6 8
 
*/