题解 | #小红购物#

以节点1为根dfs算出所有子树的连通块的数量，再对树进行一次bfs，bfs每次访问新的结点x时，将x子树和包含根结点的子树的联通块之差加入结果。

using namespace std;
int f(int a,int b)
{
    return a>b?a-b:b-a;
}
int main()
{
    int n;
    cin>>n;
    string s;
    cin>>s;
    vector<int> e[n+1];
    int u,v;
    for(int i=0;i<n-1;i++)
    {
        cin>>u>>v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    int num[n+1];
    function<void(int,int)> dfs=[&](int father,int root)
    {
        int ans=0;
        for(int i=0;i<e[root].size();i++)
        {
            int x=e[root][i];
            if(x!=father)
            {
                dfs(root,e[root][i]);
                ans+=num[x];
                if(s[root-1]==s[x-1])
                    ans--;
            }
        }
        num[root]=ans+1;
    };
    dfs(0,1);
    int root_num=num[1];
    //for(int i=1;i<=n;i++)
        //printf("%d ",num[i]);
    long long ans=0;
    queue<int> qu;
    int get[n+1];
    memset(get,-1,sizeof get);
    get[1]=1;
    qu.push(1);
    while(!qu.empty())
    {
        int x=qu.front();
        qu.pop();
        for(int i=0;i<e[x].size();i++)
        {
            int y=e[x][i];
            if(get[y]==-1)
            {
                if(s[x-1]==s[y-1])
                {
                    ans+=f(root_num-num[y]+1,num[y]);
                }
                else
                {
                    ans+=f(root_num-num[y],num[y]);
                }
                get[y]=1;
                qu.push(y);
            }
        }
    }
    cout<<ans<<endl;
    return 0;
}