2020牛客寒假算法基础集训营1 F - maki和tree （并查集，计数）

🌵 🌵 🌵

和这道题简直一模一样QWQ

思路和上面链接那道题一样不再重复，有一点就是最后计算答案时，一开始for了两层 mp[ x ]，T掉了，后来把式子转换一下发现可以优化掉一层for。

思路是这样的，假设现在和某个黑色点相连的白色连通块大小分别为a,b,c,d，按照我的两层for计算公式是这样的：ab+ac+ad+bc+bd+cd，经过转化为：a(b+c+d)+b(c+d)+cd，所以我们维护一个连通块大小余量即可，两层for变成一层QWQ

int f[MAXN],sz[MAXN],color[MAXN];//维护仅存在白色点的连通块
int find(int x) { return x==f[x]?x:f[x]=find(f[x]);}
vector<int>mp[MAXN],bk;
signed main()
{
    int n;cin>>n;
    string s;cin>>s;
    rpp(i,n)
    {
        if(s[i-1]=='W') color[i]=0,sz[i]=1,f[i]=i;
        else color[i]=1,bk.push_back(i);
    }
    rep(i,n-1)
    {
        int x,y;cin>>x>>y;
        if(color[x]==0&&color[y]==1) mp[y].push_back(x);
        if(color[x]==1&&color[y]==0) mp[x].push_back(y);
        if(color[x]==0&&color[y]==0)
        {
            int fx=find(x),fy=find(y);
            if(fy!=fx)  sz[fx]+=sz[fy], f[fy]=f[fx];
        }
    }
    ll ans=0;
    rep(i,bk.size())
    {
        int x=bk[i],cnt=0;
        rep(j,mp[x].size())  cnt+=sz[find(mp[x][j])];
        ans+=cnt;
        for(auto a:mp[x])//这里一开始锅了
        {
            int tmp=sz[find(a)];
            cnt-=tmp;ans+=tmp*cnt;
        }
    }
    cout<<ans<<endl;
    return 0;
}