淘天笔试 第三题题解

#淘天笔试# #淘天# #笔试#
考虑换根DP，先随便找个点做根，每个点只统计这个有根树下它的子树对它的贡献，也就是if(color[u] != color[v]) color[u] += color[v];
这样我们只能统计出子树的贡献，我们还需要父亲的贡献，父亲的贡献我们考虑扩展并查集，在dfs过程中顺便合并一下就能够得到每个点所在并查集的size，如果一个点和它的父亲不同颜色，那么就是相同并查集，直接用并查集size减去dp求出的子树贡献就是父亲的贡献，然后第二遍dfs，把每个点任意一种颜色都试一下，我们提前知道子树的贡献是多少，可以根据现在我们变的颜色统计一下现在的子树和，看看父亲跟他现在是不是一个颜色，是的话不加贡献，不然加上贡献，和ans取最大值即可。
代码如下：
#include
# include
#include
# include
#include
# include
#include
# include
#include
# include
#include
# include
#include
using ll = long long;
using ull = unsigned long long;
const int maxn = 1e6 + 5;
int fa[maxn],sum[maxn],siz[maxn];
std::string color;
char col[] = {'R','G','B'};
int ans  = 0;
int find(int x){
    return x == fa[x] ? x : find(fa[x]);
}
void merge(int x,int y){
    x = find(x), y = find(y);
    if(x == y) return;
    siz[x] += siz[y];
    fa[y] = x;
}
std::vector g[maxn];

void dfs(int u,int fa){
    sum[u] = 1;
    for(int i = 0;i < g[u].size();i++){
        int v = g[u][i];
        if(v == fa) continue;
        dfs(v,u);
        if(color[u] == color[v]){
            sum[u] += sum[v];
            merge(u,v);
        }
    }
}

void dfs2(int u,int fa){
    int sfa = (find(u) == find(fa)) ? siz[find(u)] - sum[u] : siz[find(fa)];
    for(int i = 0;i < 3;i++){
        char c = col[i];
        int s = 1;
        for(auto&amp; v : g[u]){
            if(v == fa) continue;
            if(c != color[v]) s += sum[v];
        }
        if(c != color[fa]) s += sfa;
        ans = std::max(ans,s);
    }
    for(auto&amp; v: g[u]){
        if(v == fa) continue;
        dfs2(v,u);
    }
}

void solve(){
    int n;
    std::cin >> n;
    std::cin >> color;
    for(int i = 1;i <= n;i++){
        siz[i] = 1;
        fa[i] = i;
    }
    for(int i = 0;i < n - 1;i++){
        int u,v;
        std::cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    dfs2(1,0);
    std::cout << ans << '\n';
}

int main() {
    int T;
    T = 1;
    while(T--) solve();
    return 0;
}