关注塔子哥学算法，全网同名。获得更多最新大厂真题+题解以及在线评测网站！

题目内容

小红拿到了一棵树，每个节点被染成了红色或者蓝色。

小红定义每条边的权值为：删除这条边时，形成的两个子树的同色连通块数量之差的绝对值。

小红想知道，所有边的权值之和是多少？

输入描述

第一行输入一个正整数 $nn$ ,代表节点的数量。

第二哈输入一个长度为 $nn$ 且仅由 $RR$ 和 $BB$ 两种字符组成的字符串。

第 $ii$ 个字符为 $RR$ 代表 $ii$ 号节点被染成红色，为 $BB$ 则被染成蓝色。

接下来的 $n-1n-\; 1$ 行，每行输入两个正整数 $uu$ 和 $vv$ ，代表节点 $uu$ 和节点 $vv$ 有一条边相连。

$1\le n\le 2000001\backslash le\; n\backslash le\; 200000$

输出描述

一个正整数，代表所有节点的权值之和。

样例

输入

4
BBRR
1 2
3 2
4 1

输出

2

题目思路

解法1:树形dp

关于树形dp,不熟悉的可以看此处学习,把基础的题目过了。

一个显然的思路是:令$d{p}_{i}dp\_i$ 代表以$ii$点为根的子树中同色联通块的个数。那么如果我们从1号点开始算，最后$d{p}_{1}dp\_1$得到的就是整棵树的同色联通块个数。现在考虑它的转移:

1.假设点 $ii$ 只有一个儿子$jj$，那么也就是只有一颗子树。

​ 那么转移很显然 : $d{p}_i=d{p}_j+[colo{r}_i\mathrm{\ne }colo{r}_j]dp\_i\; =\; dp\_j\; +\; [color\_i\; \backslash neq\; color\_j]$

两种情况:1.i,j颜色不同，连通块个数++。2.i,j颜色相等,则同色连通块得到扩展，连通块个数不变。

2.假设点$ii$ 有多个儿子$j_1,j_2,\mathrm{.}\mathrm{.}\mathrm{.},j_{k}j\_1,j\_2,...,j\_k$ , 可以将这个过程分阶段进行，即视作是将这$kk$个儿子一个一个拼接在点 $ii$ 上，所以反复进行1的转移即可。

这个过程如下图所示:

接着，我们枚举每一条边$(u\to v)(u\; \backslash rightarrow\; v)$(这个过程可以使用dfs)，断开之后原图分成两个部分.子树部分显然是$d{p}_{v}dp\_v$ , 而子树的补集是整体的同色连通块个数$d{p}_{1}dp\_1$ 减去$d{p}_{v}dp\_v$，但是注意如果$colo{r}_u=colo{r}_{v}color\_u\; =\; color\_v$ , 那么答案就需要加一(画个图就知道啦~)

所以第$ii$ 条边的贡献是:$\mathrm{\mid }(d{p}_1-d{p}_v+[colo{r}_u=colo{r}_v])-d{p}_v\mathrm{\mid }|(dp\_1\; -\; dp\_v\; +\; [color\_u\; =\; color\_v])\; -\; dp\_v|$ , 把所有的贡献加起来即可。

最后注意，由于极端情况下贡献和是$O(n^2)O(n^2)$ 的，所以记得开$longlonglong\backslash \; long$

时空复杂度:$O(n)O(n)$

解法2:dfs记忆化 + 线段树维护

群友给的思路。待补

代码+解析(基于解法1)

C++

#include<bits/stdc++.h>
using namespace std;
const int maxn = 2e5 + 5;
string a;
// 邻接矩阵 e
vector<int> e[maxn];
long long ans = 0;
// dp定义如题解所示
int dp[maxn];
void dfs (int u , int fa){
    // 如上图，在dfs之前假设最开始只有u孤零零的一个点
    dp[u] = 1;
    for (int v : e[u]){
        if (v == fa) continue;
        // 每次新增一个子树v,插到u上
        dfs(v , u);
        // 转移
        dp[u] += dp[v];
        if (a[u - 1] == a[v - 1]) dp[u]--;
    }
}

void dfs1 (int u , int fa){
    for (int v : e[u]){
        if (v == fa) continue;
        dfs1(v , u);
        // 计算边的贡献
        int x = dp[1] - dp[v];
        if (a[u - 1] == a[v - 1]) x += 1;
        ans += abs(x - dp[v]);
    }
}
int main()
{
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    cin >> a;
    for (int i = 0 ; i < n - 1; i++){
        int x , y;
        cin >> x >> y;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    // 第一遍dfs求出dp数组
    dfs(1 , -1);
    // 第二遍dfs计算每条边的贡献(也可以直接枚举邻接矩阵e)
    dfs1(1 , -1);
    cout << ans << endl;
    return 0;
}

Java - 赛时需要上快读才能拿满分

import java.util.*;
public class Main {
    static List<Integer> [] e = new List [200005];
    static String a = new String();
    static long ans = 0;
    static int [] dp = new int [200005];
    public static void dfs (int u , int fa){
        dp[u] = 1;
        for (int v : e[u]){
            if (v == fa) continue;
            dfs(v , u);
            dp[u] += dp[v];
            if (a.charAt(u - 1) == a.charAt(v - 1)) dp[u]--;
        }
    }
    public static void dfs1 (int u , int fa){
        for (int v : e[u]){
            if (v == fa) continue;
            dfs1(v , u);
            int x = dp[1] - dp[v];
            if (a.charAt(u - 1) == a.charAt(v - 1)) {
                x += 1;
            }
            ans += Math.abs(x - dp[v]);
        }
    }
    public static void main (String argc[]){
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        for (int i = 0 ; i <= n ; i++)
            e[i] = new ArrayList<>();
        a = sc.nextLine();
        a = sc.nextLine();
        for (int i = 0 ; i < n - 1 ; i++){
            int x = sc.nextInt();
            int y = sc.nextInt();
            e[x].add(y);
            e[y].add(x);
        }
        dfs(1 , -1);
        dfs1(1 , -1);
        System.out.println(ans);
        return ;
    }
}