题解 | #小红购物#
小红购物
https://ac.nowcoder.com/acm/contest/72266/A
以节点1为根dfs算出所有子树的连通块的数量,再对树进行一次bfs,bfs每次访问新的结点x时,将x子树和包含根结点的子树的联通块之差加入结果。
using namespace std;
int f(int a,int b)
{
return a>b?a-b:b-a;
}
int main()
{
int n;
cin>>n;
string s;
cin>>s;
vector<int> e[n+1];
int u,v;
for(int i=0;i<n-1;i++)
{
cin>>u>>v;
e[u].push_back(v);
e[v].push_back(u);
}
int num[n+1];
function<void(int,int)> dfs=[&](int father,int root)
{
int ans=0;
for(int i=0;i<e[root].size();i++)
{
int x=e[root][i];
if(x!=father)
{
dfs(root,e[root][i]);
ans+=num[x];
if(s[root-1]==s[x-1])
ans--;
}
}
num[root]=ans+1;
};
dfs(0,1);
int root_num=num[1];
//for(int i=1;i<=n;i++)
//printf("%d ",num[i]);
long long ans=0;
queue<int> qu;
int get[n+1];
memset(get,-1,sizeof get);
get[1]=1;
qu.push(1);
while(!qu.empty())
{
int x=qu.front();
qu.pop();
for(int i=0;i<e[x].size();i++)
{
int y=e[x][i];
if(get[y]==-1)
{
if(s[x-1]==s[y-1])
{
ans+=f(root_num-num[y]+1,num[y]);
}
else
{
ans+=f(root_num-num[y],num[y]);
}
get[y]=1;
qu.push(y);
}
}
}
cout<<ans<<endl;
return 0;
}