2020牛客寒假算法基础集训营1

这场比较基础吧

maki和tree

直接点分治也可以吧。

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#include <bits/stdc++.h>
#define ll long long
#define lld unsigned long long
#define inf 0x3f3f3f3f
#define mod 1000000007
using namespace std;
vector<int> g[100010];
int mx=inf;
bool vis[100010];
int son[100010],sz[100010];
ll res[100010],q[10];
int sum,rt;
//ll ans1,ans2,ans3;
ll ans=0;
char black[100010];
void getroot(int x,int fa)
{
    son[x]=0;sz[x]=1;
    for(auto u: g[x])
    {
        int v=u;
        if(vis[v]||v==fa) continue;
        getroot(v,x);
        sz[x]+=sz[v];
        if(sz[v]>son[x]) son[x]=sz[v];
    }
    if(sum-sz[x]>son[x]) son[x]=sum-sz[x];
    if(mx>son[x]){
        mx=son[x];
        rt=x;
    }
    return;
}
void getdis(int x,int fa,ll len)
{
    res[++res[0]]=len;
    for(auto v: g[x])
    {
        if(vis[v]||v==fa)continue;
        getdis(v,x,len+(black[v]=='B'));
    }
}
void calc(int x)
{
    for(auto v: g[x])
    {
        if(vis[v]) continue;
        res[0]=0;
        getdis(v,x,(black[v]=='B'));
        ll cnt1=0;ll cnt0=0;
        for(int i=1;i<=res[0];i++)
        {
            if(res[i]==1) cnt1++;
            if(res[i]==0) cnt0++;
        }
        if(black[x]=='B')
        {
            ans+=cnt0*q[0]+cnt0;
        }
        else
        {
            ans+=cnt0*q[1]+q[0]*cnt1+cnt1;
        }
        for(int i=1;i<=res[0];i++)//将本子树的单链保存
        {
            if(res[i]==1) q[1]++;
            if(res[i]==0) q[0]++;
        }
    }
    for(int i=0;i<3;i++)
    {
        q[i]=0;
    }
}
void solve(int u)
{
    rt=0;
    mx=inf;
    getroot(u,0);//找根
    vis[rt]=1;
    calc(rt);//统计
    for(auto v: g[rt])
    {
        if(vis[v]) continue;
        sum=sz[v];
        solve(v);
    }
}
int main()
{
    int n;
    scanf("%d",&n);
    scanf("%s",black+1);
    for(int i=1;i<=n-1;i++)
    {
        int u,v;
        scanf("%d %d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    sum=n;
    solve(1);
    printf("%lld\n",ans);
    return 0;
}