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

题目链接:https://ac.nowcoder.com/acm/contest/3002/F

题解思路:

思路统计连通块白点的个数(可用并查集，或者搜索)

• 搜索: https://ac.nowcoder.com/acm/contest/view-submission?submissionId=42763870
• 并查集: https://blog.csdn.net/qq_42815188/article/details/104220706
  

代码

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;
vector<vector<int> > mp;
int vis[MAXN], num[MAXN];
vector<int> tmp;
int bfs(int now) {
    queue<int> q;
    q.push(now);
    tmp.push_back(now);
    vis[now] = 1;
    int cnt = 1;        // now 也算一个
    while(!q.empty()) {
        int k = q.front();
        q.pop();
        int vsize = mp[k].size();
        for(int i = 0; i < vsize; i++){
            if(vis[mp[k][i]] == 0) {
                q.push(mp[k][i]);
                vis[mp[k][i]] = 1;
                tmp.push_back(mp[k][i]);
                cnt++;
            }
        }
    }

    return cnt;
}
int main()
{
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    mp.resize(n + 5);
    string s;
    cin >> s;
    fill(vis, vis + MAXN, 0);
    int len = s.length();
    for(int i = 0; i < len; i++){
        if(s[i] == 'B'){
            // 黑色
            vis[i + 1] = 1;
        }
    }
    for(int i = 1; i < n; i++){
        int a, b;
        cin >> a >> b;
        mp[a].push_back(b);
        mp[b].push_back(a);
    }
    long long ans = 0;
    for(int i = 1; i <= n; i++){
        // 如果是黑色
        if(vis[i] == 1) {
            int now = 0;
            int sum = 1;    // 分两种情况， 1. 一个端点是黑色， 2. 两点都是白色
            int vsize = mp[i].size();
            for(int j = 0; j < vsize; j++){
                // 该点是白色
                if(vis[mp[i][j]] == 0) {
                    num[now++] = bfs(mp[i][j]);
                    sum += num[now - 1];
                }
            }
            // 去除标记
            while(!tmp.empty()) {
                vis[tmp.back()] = 0;
                tmp.pop_back();
            }
            // 统计
            for(int j = 0; j < now; j++){
                ans += num[j] * (sum - num[j]);     // 1. num[j] * (sum - num[j] - 1) 2. num[j] * 1
                sum = sum - num[j];
            }
        }
    }

    cout << ans << endl;
    return 0;
}