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

https://ac.nowcoder.com/acm/contest/3002/F

题意

　　这个树有 个顶点， 条边。每个顶点被染成了白色或者黑色。
　　取两个不同的点，它们的简单路径上有且仅有一个黑色点的取法有多少？
　　注：
　　　　　　①树上两点简单路径指连接两点的最短路。
　　　　　　② 和 的取法视为同一种。

题解

　　并查集+计数。

　　只经过一个黑色点的路径无非有两种，① 两端为白点，中间有一个黑点；② 任意一端是黑点。

　　预处理，将每个白点的最大连通块处理出来，记录其大小。

　　若一个黑点与 k 个白点相连，令 f(i) 为第 i 个白点的连通块大小（）。

　　则满足 ① 情况的数量为：

　　满足 ② 情况的数量为：

代码

#include<bits/stdc++.h>
using namespace std;
const int maxn=100005;
int N,f[maxn],link[maxn];
string is_wb;
vector<int> edge[maxn];
int find(int p);
void merge(int u,int v);
long long cal();
int main()
{
   {
      fill(f,f+maxn,-1);
   }
   int i,u,v;
   scanf("%d",&N);
   cin>>is_wb;
   for(i=1;i<N;i++)
   {
      scanf("%d%d",&u,&v);
      edge[u].push_back(v);
      edge[v].push_back(u);
      if(is_wb[u-1]=='W' && is_wb[v-1]=='W')
         merge(u,v);
   }
   printf("%lld",cal());
   system("pause");
   return 0;
}
int find(int p)
{
   return f[p]==-1?p:f[p]=find(f[p]);
}
void merge(int u,int v)
{
   int x,y;
   x=find(u);
   y=find(v);
   if(x!=y)
   {
      f[x]=y;
      link[y]+=link[x]+1;
   }
}
long long cal()
{
   int i,j;
   long long ans=0;
   for(i=1;i<=N;i++)
   {
      vector<int> v;
      vector<int> pre_v;
      v.push_back(0);
      pre_v.push_back(0);
      if(is_wb[i-1]=='W') continue;
      for(j=0;j<edge[i].size();j++)
      {
         if(is_wb[edge[i][j]-1]=='B') continue;
         v.push_back(link[find(edge[i][j])]+1);
         pre_v.push_back(link[find(edge[i][j])]+1+pre_v.back());
      }
      for(j=1;j<v.size();j++)
         ans+=v[j]*(pre_v.back()-pre_v[j]);
      ans+=pre_v.back();
   }
   return ans;
}