游游的三色树（2022携程秋招第一批）：题解

题目链接：https://www.nowcoder.com/questionTerminal/92756623de4146d1a573907478947987

一开始以为是换根dp，结果其实是个比较简单的树上计数题。

假设树的根是id=1的那个节点，定义siz[i][j][k]表示第i个节点，考虑其子树上每个节点，颜色为j的节点之和，k=0代表其子树，k=1代表除了其子树以外的区域，siz[i][j][0]用dfs直接求，siz[i][j][1]实际就是拿根节点（即整棵树）的siz减去siz[i][j][0]，最后看是不是都>=1即可。

颜色3种，k有2种，dfs本身是O(n)的，所以总体是O(6*n)=O(n)的，就是可能常数有点大。

#include <iostream>
#include <string>
#include<vector>
#include<cstring>
using namespace std;
using ll = long long;
string s;
int n;
vector<int>g[100005];
ll siz[100005][3][2];
void dfs(int u, int fa){
    for(auto v : g[u]){
        if(v == fa){
            continue;
        }
        dfs(v,u);
        siz[u][0][0] += siz[v][0][0], siz[u][1][0] += siz[v][1][0], siz[u][2][0] += siz[v][2][0];
    }
}
void dfs2(int u, int fa){
    if(u != 1){
        siz[u][0][1] = siz[1][0][0]-siz[u][0][0];
        siz[u][1][1] = siz[1][1][0]-siz[u][1][0];
        siz[u][2][1] = siz[1][2][0]-siz[u][2][0];
    }
    for(auto v : g[u]){
        if(v == fa){
            continue;
        }
        dfs2(v,u);
    }
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    cin >> n;
    cin >> s;
    memset(siz,0,sizeof siz);
    for(int i = 1;i <= n;++i){
        if(s[i-1] == 'r'){
            siz[i][0][0]++;
        }
        else if(s[i-1] == 'g'){
            siz[i][1][0]++;
        }
        else{
            siz[i][2][0]++;
        }
    }
    for(int i = 1,u,v;i < n;++i){
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    dfs2(1,0);
    ll ans = 0;
    for(int i = 2;i <= n;++i){
        if(siz[i][0][0] >= 1 && siz[i][1][0] >= 1 && siz[i][2][0] >= 1 && siz[i][0][1] >= 1 && siz[i][1][1] >= 1 && siz[i][2][1] >= 1){
            ans++;
        }
    }
    cout << ans << "\n";
    return 0;
}
// 64 位输出请用 printf("%lld")