原题地址:https://ac.nowcoder.com/acm/contest/3002/F
解题思路
深搜统计每个黑点相连的白点的个数,一个黑点的ans为所有与它相连的白点个数加上,两两两条线之间的乘积,把每个黑点的ans相加就是答案
AC代码
#include<iostream>
#include<vector>
#include<map>
#include<set>
using namespace std;
typedef long long ll;
const int maxn = 100000 + 10;
char colors[maxn];
int n;
ll ans = 0;
struct point
{
vector<int> next;
char color;
vector<int> count;
};
point tree[maxn];
int dfs(int cur, int last)
{//统计每一个黑点相邻的白点数量
int res = 1;
for (int i = 0; i < (int)tree[cur].next.size(); i++)
{
int nextP = tree[cur].next[i];
if (nextP == last) continue;
if (tree[nextP].color == 'W')
{
res += dfs(nextP, cur);
}
}
return res;
}
int main()
{
scanf("%d", &n);
scanf("%s", colors);
for (int i = 1; i <= n; i++)
{
tree[i].color = colors[i - 1];
}
for (int i = 0; i < n - 1; i++)
{
int a, b;
scanf("%d%d", &a, &b);
tree[a].next.push_back(b);
tree[b].next.push_back(a);
}
for (int i = 1; i <= n; i++)
{
if (tree[i].color == 'B')
{
int l = tree[i].next.size();
for (int j = 0; j < l; j++)
{
int nextP = tree[i].next[j];
if (tree[nextP].color == 'W')
{
tree[i].count.push_back(dfs(nextP, i));
}
}
vector<int> count = tree[i].count;
for (int j = 0; j < (int)count.size(); j++)
{
ans += count[j];
for (int k = j + 1; k < (int)count.size(); k++)
{
ans += (count[j] * count[k]);
}
}
}
}
cout << ans;
return 0;
}