77 lines
2.1 KiB
C++
77 lines
2.1 KiB
C++
#include <bits/stdc++.h>
|
||
|
||
using namespace std;
|
||
class Solution {
|
||
public:
|
||
int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
|
||
int n=graph.size();
|
||
vector<int> sz(n,1); //n号节点如果是祖先，他的孩子个数
|
||
vector<int> v_num(n,0);//n号节点如果是祖先，他的孩子中initial的个数
|
||
for (int i: initial){
|
||
v_num[i]=1;
|
||
}
|
||
//并查集核心模板代码
|
||
//并查集祖先数组定义
|
||
vector<int> pa(n);
|
||
//并查集祖先数组初始化
|
||
for (int i=0;i<n;i++) {
|
||
pa[i]=i;
|
||
}
|
||
//并查集查找祖先函数
|
||
function<int (int)> query = [&](int i){
|
||
if (pa[i]!=i)
|
||
pa[i]=query(pa[i]);
|
||
return pa[i];
|
||
};
|
||
//并查集合并祖先函数
|
||
auto link = [&](int i,int j){
|
||
int li=query(i);
|
||
int lj=query(j);
|
||
if (li != lj){
|
||
pa[li]=lj;
|
||
//其他附加操作
|
||
sz[lj]+=sz[li];
|
||
sz[li]=0;
|
||
v_num[lj]+=v_num[li];
|
||
v_num[li]=0;
|
||
}
|
||
return;
|
||
};
|
||
//
|
||
//建立并查集，合并祖先
|
||
for (int i=0;i<n;i++){
|
||
for (int j=0;j<n;j++){
|
||
if (graph[i][j]==1)
|
||
link(i,j);
|
||
}
|
||
}
|
||
int res=n;
|
||
int res_multi=n;
|
||
int max_nodes=0;
|
||
bool have_one=false;
|
||
for (int i:initial){
|
||
int li=query(i);
|
||
if (v_num[li]==1){
|
||
if (sz[li]>max_nodes){
|
||
res=i;
|
||
max_nodes=sz[li];
|
||
}else if (sz[li]==max_nodes){
|
||
res=min(res,i);
|
||
}
|
||
have_one=true;
|
||
} else {
|
||
res_multi=min(res_multi,i);
|
||
}
|
||
}
|
||
return have_one?res:res_multi;
|
||
}
|
||
};
|
||
int main()
|
||
{
|
||
Solution sol;
|
||
vector<vector<int>> graph={{1,0,0},{0,1,0},{0,0,1}};
|
||
vector<int> initial={0,2};
|
||
cout << sol.minMalwareSpread(graph,initial) << endl;
|
||
return 0;
|
||
}
|