牛客-maki和tree（dfs）

qaq好久没写博客了，假期使人不想学习，，，

题目链接： 传送门

题目描述
有一天，maki拿到了一颗树。所谓树，即没有自环、重边和回路的无向连通图。
这个树有 个顶点， 条边。每个顶点被染成了白色或者黑色。
maki想知道，取两个不同的点，它们的简单路径上有且仅有一个黑色点的取法有多少？
注：
①树上两点简单路径指连接两点的最短路。
② 和 的取法视为同一种。
输入描述:
第一行一个正整数 。代表顶点数量。
第二行是一个仅由字符’B’和’W’组成的字符串。第 个字符是B代表第 个点是黑色，W代表第 个点是白色。
接下来的 行，每行两个正整数 ， ，代表 点和 点有一条边相连
输出描述:
一个正整数，表示只经过一个黑色点的路径数量。
示例1
输入
复制
3
WBW
1 2
2 3
输出
复制
3
说明
树表示如下：

其中只有2号是黑色点。
<1,2>、<2,3>、<1,3>三种取法都只经过一个黑色点。

思路：开始没想到是什么算法，就暴力大法好，直接枚举每个黑色点相连的所有分支，找到每个分支的可选点有多少个，然后两两相乘再相加，，交上去T了，觉得暴力大法不行，然后看别的题去了，，之后lc学长跟我说他n*n伪暴力1A，我？？？后来发现我是直接两重循环相乘相加的，其实可以搞个前缀和优化掉一重循环，，就过了，，我为啥没想到呢？？，，可能心里本来就在想别的题，，我fo了

代码：

#include
#include
#include
#include
#include
#include
using namespace std;
#define inf 0x3f3f3f3f
const int N = 2e5+10;
typedef long long ll;
char s[N];
int e[N],ne[N],h[N],idx;
vector<int> tem;
ll total=0;
void add(int a,int b)
{
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
int dfs(int x,int pre)
{
    int sum=1,xx;
    for(int i=h[x];i!=-1;i=ne[i])
    {
        int j=e[i];
        if(j==pre) continue;
        if(s[j]=='W')
            xx=dfs(j,x),sum+=xx;;
        if(pre==-1&&s[j]=='W')
            tem.push_back(xx),total+=xx;

    }
    return sum;
}
int main()
{
    ll ans=0;
    int n;
    cin>>n;
    memset(h,-1,sizeof h);
    scanf("%s",s+1);
    for(int i=0;i<n-1;i++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b);
        add(b,a);
    }
    for(int i=1;i<=n;i++)
    {
        if(s[i]=='B')
        {
            tem.clear();
            total=0;
            ans+=dfs(i,-1)-1;
            //cout<<i<<" "<<tem.size()<<" "<<total<<" "<<ans<<endl;
            for(int i=0;i<tem.size();i++)
            {
                ans+=(ll)tem[i]*(total-tem[i]);
                total-=tem[i];
            }
        }
    }
    cout<<ans<<endl;
    return 0;
}