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;
|
|||
|
}
|