题解 |牛客周赛 Round 25D 小红树

D-小红树

首先以u为根，通过dfs统计以u为根的所有子树的同色连通块数

后用dfs 计算abs(差值)累加于res中

连通块：

连通块的数量n与树上的节点数n相当

同色连通块：

n - （相连同色的节点 - 1）;

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const int N = 2e5 + 10;
int n;          //n为节点总数
int f[N];       //f[]储存连通块数量。以f[i]为根，向下加上自身的连通块数
string s;       //存储颜色
vector<int>g[N];//用于储存树
ll res;         //统计总数

//通过深搜遍历树，统计以u为根的连通块的数目
void dfs_1(int u, int father)
{
    f[u] = 1;//赋值 根u 连通块初始值为 1(自身)
    for (int x:g[u])//遍历g[u]为根向下的节点
    {
        if (x == father) continue;//当遍历到指向的根方向时，跳过
        dfs_1(x, u);
        f[u] += f[x];   //统计连通块数量
        if (s[u] == s[x]) f[u] --;//题目要求统计同色连通块
        //当遍历到同色连通块时  f[u] - 1
    }
}

//计算两个子树同色连通块的差值，并统计存入res中
void dfs_2(int u, int father)
{
    for (int x:g[u])//遍历g[u]为根向下的节点
    {
        if (x == father) continue;//当遍历到指向的根方向时，跳过
        dfs_2(x, u);
        //f[x]为以x为根的子树的连通块数量
        //v 为另一个子树的连通块的数量
        int v = f[1] - f[x];
        if (s[u] == s[x]) v ++;//若原本为同色，在分开始必然增加同色连通块的数量
        res += abs(v - f[x]);//统计，存入res中
    }
}

int main()
{
    int a, b;
    ios::sync_with_stdio(false);
    cin >> n >> s;
    s = '$' + s;//使颜色与f[]的下标相统一
    n --;
    while(n --)
    {
        //输入结点，建立连接结点的线
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    dfs_1(1, -1);
    dfs_2(1, -1);
    cout << res << endl;

    return 0;
}

代码来源---- 牛客：2333_