2020牛客寒假算法基础训练营Day1 - F maki和tree

传送门：Day1 - F

题目描述

思路分析

将相邻的白色节点连通，记录每个连通块内点的数量，再枚举每一个黑点。
如图：

这张图里的 $ans=(1+2+3)+(1×2)+(1×3)+(2×3)$

样例输入

3
WBW
1 2
2 3

样例输出

3

AC代码

#include <iostream>
#include <string>
#include <cstring>
#include <cstdio>
#include <map>
#include <set>
#include <queue>
#include <stack>
#include <algorithm>
#include <cstdlib>
#include <cmath>
#include <vector>
#include <iomanip>
#if __cplusplus >= 201103L
#include <unordered_map>
#include <unordered_set>
#endif
#define fi first
#define se second
#define ll long long
#define pb push_back
#define vi vector<int>
#define lowbit(x) x&(-x)
#define pii pair<int,int>
#define umap unordered_map
#define uset unordered_set
#define all(x) x.begin(),x.end()
#define mem(a,b) memset(a,b,sizeof(a))
#define IOS ios_base::sync_with_stdio(0); cin.tie(0);cout.tie(0);
 
using namespace std;
 
const int INF = 0x3f3f3f3f;
const int N=1e5+10;
 
ll f[N],size[N];
vector<int>e[N];
char s[N];
 
int find(int x) {
    if(f[x]==x)
        return x;
    else
        return f[x]=find(f[x]);
}
 
void merge(int u,int v) {
    int t1=find(u);
    int t2=find(v);
    if(t1!=t2) {
        f[t2]=t1; 
        size[t1]+=size[t2];
    }
}
 
 
int main() {
    IOS;
#ifdef xiaofan
    freopen("in.txt","r",stdin);
#endif
 
    int n;
    cin>>n;
    cin>>s+1;
    for(int i=1; i<=n; i++) {
        if(s[i]=='W') size[i]++;
        f[i]=i;
    }
    for(int i=1; i<=n-1; i++) {
        int u,v;
        cin>>u>>v;
        e[u].pb(v);
        e[v].pb(u);
        if(s[u]==s[v] && s[u]=='W')
            merge(u,v);
    }
    ll ans=0;
    for(int i=1; i<=n; i++) {
        if(s[i]=='B') {
            ll sum=0;
            for(auto v:e[i]) sum+=size[find(v)];
            ans+=sum;
            for(auto v:e[i]) {
                sum-=size[find(v)];
                ans+=sum*size[find(v)];
            }
        }
    }
    cout<<ans<<endl;
 
 
 
 
 
    return 0;
}