Problem_0928_minMalwareSpread.cc

#include <algorithm>
#include <numeric>
#include <vector>

#include "UnitTest.h"

using namespace std;

// TODO: figure it out.
// @sa https://www.bilibili.com/video/BV1Ny4y1F71J Code04
class Solution
{
 private:
  // [3,6,103]
  // virus[3] = true;
  // virus[103] = true;
  // 方便查询
  vector<bool> virus;
  // 每个源头点删掉的话，能拯救多少点的数据
  vector<int> cnts;

  // 集合的标签 : 集合的感染点是什么点
  // a : 代表点，整个集合源头是 infect[a]
  // infect[a] == -1，目前这个集合没有发现源头
  // infect[a] >= 0，目前这个集合源头是 infect[a]
  // infect[a] == -2，目前这个集合源头不止一个，已经无法拯救了!
  vector<int> infect;

  // 并查集固有信息
  vector<int> father;

  // 集合的标签 : 集合的大小是多少
  vector<int> size;

  // 集合一定只放普通点，源头点根本不参与集合，也不是元素！

  void build(int n, vector<int>& initial)
  {
    virus.resize(n, false);
    cnts.resize(n, 0);
    infect.resize(n, -1);
    size.resize(n, 1);
    father.resize(n);
    std::iota(father.begin(), father.end(), 0);
    for (auto& i : initial)
    {
      virus[i] = true;
    }
  }

  int find(int i)
  {
    if (i != father[i])
    {
      father[i] = find(father[i]);
    }
    return father[i];
  }

  void unions(int x, int y)
  {
    int fx = find(x);
    int fy = find(y);
    if (fx != fy)
    {
      father[fx] = fy;
      size[fy] += size[fx];
    }
  }

 public:
  int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial)
  {
    int n = graph.size();
    build(n, initial);
    // 不是病毒的点，普通点合并！
    for (int i = 0; i < n; i++)
    {
      for (int j = 0; j < n; j++)
      {
        if (graph[i][j] == 1 && !virus[i] && !virus[j])
        {
          unions(i, j);
        }
      }
    }
    // 病毒周围的普通点(集合 )去设置源头！
    for (int sick : initial)
    {
      for (int neighbor = 0; neighbor < n; neighbor++)
      {
        if (sick != neighbor && !virus[neighbor] && graph[sick][neighbor] == 1)
        {
          int fn = find(neighbor);
          if (infect[fn] == -1)
          {
            infect[fn] = sick;
          }
          else if (infect[fn] != -2 && infect[fn] != sick)
          {
            infect[fn] = -2;
          }
        }
      }
    }
    // 统计拯救数据
    for (int i = 0; i < n; i++)
    {
      // 不是代表点，不看
      if (i == find(i) && infect[i] >= 0)
      {
        cnts[infect[i]] += size[i];
      }
    }
    std::sort(initial.begin(), initial.end());
    int ans = initial[0];
    int max = cnts[ans];
    for (int i : initial)
    {
      if (cnts[i] > max)
      {
        ans = i;
        max = cnts[i];
      }
    }
    return ans;
  }
};

void test()
{
  Solution s;
  vector<vector<int>> g1 = {{1, 1, 0}, {1, 1, 0}, {0, 0, 1}};
  vector<int> ini = {0, 1};
  vector<vector<int>> g2 = {{1, 1, 0}, {1, 1, 1}, {0, 1, 1}};
  vector<vector<int>> g3 = {{1, 1, 0, 0}, {1, 1, 1, 0}, {0, 1, 1, 1}, {0, 0, 1, 1}};
  EXPECT_EQ_INT(0, s.minMalwareSpread(g1, ini));
  EXPECT_EQ_INT(1, s.minMalwareSpread(g2, ini));
  EXPECT_EQ_INT(1, s.minMalwareSpread(g3, ini));
  EXPECT_SUMMARY;
}

int main()
{
  test();
  return 0;
}