2020牛客寒假算法基础集训营1 F maki和tree(图论+dfs)

题目链接
思路：以每个黑点作为根，统计一下白色点的连通块大小，计数即可。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 2e5 + 10;
char s[N] ;
ll ans , n, a[N];
queue<int> q ;
int tot , head[N] ;
struct Edge
{
    int v , next ;
} edge[N << 1] ;
void add_edge(int u , int v)
{
    edge[tot].v = v ;
    edge[tot].next = head[u] ;
    head[u] = tot ++ ;
}
void dfs(int f , int u , int id)
{
    ans ++ ;
    q.push(id) ;
    for(int i = head[u] ; i != -1 ; i = edge[i].next)
    {
        int v = edge[i].v ;
        if(v == f || s[v] == 'B')  continue ;
        dfs(u , v , id) ;
    }  
}
void solve()
{
    int cot = 0 , now = 0 ;
    ll sum = q.size() ;
    ll cnt = 0 ;
    while(!q.empty())
    {
       int x = q.front() ;
       q.pop() ;   
       if(x != now) 
         now = x , a[++ cot] = 1 ;
       else  a[cot] ++ ;
    }
    
    for(int i=1;i<=cot;++i)
    cnt += a[i] * (sum - a[i]) ;
    ans += cnt / 2 ;
}
int main()
{
    scanf("%d",&n);
    scanf("%s",s+1);
    tot = 0;
    memset(head,-1,sizeof(head));
    for(int i=1;i<n;++i)
    {
        int u , v ;
        cin >> u >> v;
        add_edge(u , v) ;
        add_edge(v , u) ;
    }
    for(int i=1;i<=n;++i)
    {
        if(s[i] == 'W')  continue ;
        while(!q.empty())  q.pop() ;
        for(int j = head[i] ; j != -1 ; j = edge[j].next)
        {
           int v = edge[j].v ;
           if(s[v] == 'B')  continue ;
           dfs(i , v , v) ;
        }
        solve() ;
    }
   printf("%lld\n",ans);
}