字节9.3秋招研发笔试 【后端方向】第三题

题目

小红拿到了一个无向图，初始每人节点是白色，其中有若干个节点被染成了红色。小红想知道，若将 i 号节点染成红色，当前的红色连块的数量是多少? 你需要回答i∈[1,n] 的答案。

定义，若干节点组成一个红色连通块，当且仅当它们都是红色节点，且在该图上可以通过无向边互相到达，这些可以连通的节点构成的最大集合为一个连通块。

代码

考察：并查集，建图

#include <iostream>
#include <bits/stdc++.h>
using namespace std;
/*
    保存所有边
    当边的两端节点都为红色，合并连通块
    遍历节点：
        若为红，输出当前连通块数量all；
        若为白，遍历其连接的红节点，得出其连接的连通块数量x，cnt_i = all - x + 1
*/
vector<vector<bool>> edges;
string colors; // 记录颜色
int cnt= 0;  // 连通块数量
class UnionFind {
    vector<int> parents;
public:
    UnionFind(int n) {
        parents.resize(n + 1);
        for(int i = 1; i <= n; i++) {
            parents[i] = i;
        }
    }
    void Union(int x, int y) {
        int fx = find(x);
        int fy = find(y);
        if(fx != fy) {
            parents[fy] = fx;
        }
        return;
    }
    int find(int x) {
        if(parents[x] != x) {
            parents[x] = find(parents[x]);
        }
        return parents[x];
    }
    int cnt_red() {
        int num = 0;
        for(int i = 1; i < parents.size(); i++) {
            if(parents[i] == i && colors[i - 1] == 'R')
                num++;
        }
        return num;
    }
};

int main() {
    int n, m;
    cin >> n >> m;
    UnionFind* uf = new UnionFind(n);
    cin >> colors;
    edges.resize(n + 1);
    for(int i = 1; i <= n; i++) {
        edges[i].resize(n + 1);
    }
    int u, v;
    for(int i = 0; i < m; i++) {
        cin >> u >> v;
        edges[u][v] = true;
        edges[v][u] = true;
        if(colors[u - 1] == 'R' && colors[v - 1] == 'R') {
            uf->Union(u, v);
        }
    }
    cnt = uf->cnt_red();
    /*
        遍历节点：
        若为红，输出当前连通块数量all；
        若为白，遍历其连接的红节点，得出其连接的连通块数量x，cnt_i = all - x + 1
    */
   for(int i = 1; i <= n; i++) {
        if(colors[i - 1] == 'R') {
            cout << cnt;
        }
        else {
            unordered_set<int>st; // 存储i点连接的连通块
            for(int j = 1; j <= n; j++) {
                if(edges[i][j] == true && colors[j - 1] == 'R'){
                    st.insert(uf->find(j));
                }
            }
            cout << cnt - st.size() + 1;
        }
        if(i != n) cout << endl;
   }
    return 0;
}
/*
4 4
WRWR
1 2
2 3
3 4
1 4

1
2
1
2
*/