fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. class Solution {
  5. public:
  6.     vector<vector<int>> adj;
  7.     vector<int> comp, compSize;
  8.  
  9.     void dfs(int u, int id) {
  10.         comp[u] = id;
  11.         compSize[id]++;
  12.  
  13.         for (int v : adj[u]) {
  14.             if (comp[v] == -1)
  15.                 dfs(v, id);
  16.         }
  17.     }
  18.  
  19.     int minMalwareSpread(int n,
  20.                          vector<int>& from,
  21.                          vector<int>& to,
  22.                          vector<int>& malware) {
  23.  
  24.         adj.assign(n + 1, {});
  25.  
  26.         for (int i = 0; i < from.size(); i++) {
  27.             adj[from[i]].push_back(to[i]);
  28.             adj[to[i]].push_back(from[i]);
  29.         }
  30.  
  31.         comp.assign(n + 1, -1);
  32.         compSize.assign(n + 1, 0);
  33.  
  34.         int id = 0;
  35.  
  36.         // Find connected components
  37.         for (int i = 1; i <= n; i++) {
  38.             if (comp[i] == -1) {
  39.                 dfs(i, id);
  40.                 id++;
  41.             }
  42.         }
  43.  
  44.         vector<int> infectedCnt(id, 0);
  45.  
  46.         // Count infected nodes in each component
  47.         for (int i = 1; i <= n; i++) {
  48.             if (malware[i] == 1)
  49.                 infectedCnt[comp[i]]++;
  50.         }
  51.  
  52.         int ans = -1;
  53.         int maxSaved = -1;
  54.  
  55.         for (int i = 1; i <= n; i++) {
  56.  
  57.             if (malware[i] == 0)
  58.                 continue;
  59.  
  60.             int c = comp[i];
  61.  
  62.             if (infectedCnt[c] == 1) {
  63.  
  64.                 if (compSize[c] > maxSaved) {
  65.                     maxSaved = compSize[c];
  66.                     ans = i;
  67.                 }
  68.                 else if (compSize[c] == maxSaved && i < ans) {
  69.                     ans = i;
  70.                 }
  71.             }
  72.         }
  73.  
  74.         // If no component has exactly one infected node
  75.         if (ans == -1) {
  76.             for (int i = 1; i <= n; i++) {
  77.                 if (malware[i] == 1)
  78.                     return i;
  79.             }
  80.         }
  81.  
  82.         return ans;
  83.     }
  84. };
  85.  
  86. int main() {
  87.  
  88.     int g_nodes, g_edges;
  89.     cin >> g_nodes >> g_edges;
  90.  
  91.     vector<int> g_from(g_edges), g_to(g_edges);
  92.  
  93.     for (int i = 0; i < g_edges; i++)
  94.         cin >> g_from[i];
  95.  
  96.     for (int i = 0; i < g_edges; i++)
  97.         cin >> g_to[i];
  98.  
  99.     vector<int> malware(g_nodes + 1);
  100.  
  101.     for (int i = 1; i <= g_nodes; i++)
  102.         cin >> malware[i];
  103.  
  104.     Solution obj;
  105.  
  106.     cout << obj.minMalwareSpread(g_nodes, g_from, g_to, malware);
  107.  
  108.     return 0;
  109. }
Success #stdin #stdout 0s 5324KB
comments (?)
stdin
copy 
9 5
1 2 4 6 7
2 3 5 7 8
0 0 1 0 1  0 0 0 0
stdout
copy 
3
https://ideone.com/5g8T41
language:
C++ (gcc 8.3)
created:
2 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!