Problem_0924_minMalwareSpread.cc

#include <queue>
#include <unordered_map>
#include <vector>

using namespace std;

class Solution
{
 public:
  int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial)
  {
    int n = graph.size();
    vector<int> ids(n);
    unordered_map<int, int> id_to_size;
    int id = 0;
    for (int i = 0; i < n; i++)
    {
      // 节点没有被访问过
      if (!ids[i])
      {
        id++;
        int size = 1;
        queue<int> q;
        q.push(i);
        while (!q.empty())
        {
          int u = q.front();
          q.pop();
          for (int v = 0; v < n; v++)
          {
            if (!ids[v] && graph[u][v] == 1)
            {
              size++;
              q.push(v);
              // 记录节点 v 在哪个连通分量中
              ids[v] = id;
            }
          }
        }
        // 记录这个连通分量的节点数
        id_to_size[id] = size;
      }
    }
    unordered_map<int, int> id_to_initials;
    for (int u : initial)
    {
      // 连通分量有多少个被感染的节点
      id_to_initials[ids[u]]++;
    }
    int ans = n + 1;
    int ans_removed = 0;
    for (int u : initial)
    {
      // 1. 如果连通分量只有一个被感染的点，那么把这个感染的点变成未感染，可以让整个连通分量不被感染
      // 2. 如果连通分量不止一个被感染的点，无效
      int removed = (id_to_initials[ids[u]] == 1) ? id_to_size[ids[u]] : 0;
      if (removed > ans_removed || (removed == ans_removed && u < ans))
      {
        ans = u;
        ans_removed = removed;
      }
    }
    return ans;
  }
};