【图论】【并集查找】【C++算法】928. 尽量减少恶意软件的传播 II

class CUnionFind
{
public:
  CUnionFind(int iSize) :m_vNodeToRegion(iSize)
  {
    for (int i = 0; i < iSize; i++)
    {
      m_vNodeToRegion[i] = i;
    }
    m_iConnetRegionCount = iSize;
  }
  int GetConnectRegionIndex(int iNode)
  {
    int& iConnectNO = m_vNodeToRegion[iNode];
    if (iNode == iConnectNO)
    {
      return iNode;
    }
    return iConnectNO = GetConnectRegionIndex(iConnectNO);
  }
  void Union(int iNode1, int iNode2)
  {
    const int iConnectNO1 = GetConnectRegionIndex(iNode1);
    const int iConnectNO2 = GetConnectRegionIndex(iNode2);
    if (iConnectNO1 == iConnectNO2)
    {
      return;
    }
    m_iConnetRegionCount--;
    if (iConnectNO1 > iConnectNO2)
    {
      UnionConnect(iConnectNO1, iConnectNO2);
    }
    else
    {
      UnionConnect(iConnectNO2, iConnectNO1);
    }
  }
  bool IsConnect(int iNode1, int iNode2)
  {
    return GetConnectRegionIndex(iNode1) == GetConnectRegionIndex(iNode2);
  }
  int GetConnetRegionCount()const
  {
    return m_iConnetRegionCount;
  }
  vector<int> GetNodeCountOfRegion()//各联通区域的节点数量
  {
    const int iNodeSize = m_vNodeToRegion.size();
    vector<int> vRet(iNodeSize);
    for (int i = 0; i < iNodeSize; i++)
    {
      vRet[GetConnectRegionIndex(i)]++;
    }
    return vRet;
  }
  std::unordered_map<int, vector<int>> GetNodeOfRegion()
  {
    std::unordered_map<int, vector<int>> ret;
    const int iNodeSize = m_vNodeToRegion.size();
    for (int i = 0; i < iNodeSize; i++)
    {
      ret[GetConnectRegionIndex(i)].emplace_back(i);
    }
    return ret;
  }
private:
  void UnionConnect(int iFrom, int iTo)
  {
    m_vNodeToRegion[iFrom] = iTo;
  }
  vector<int> m_vNodeToRegion;//各点所在联通区域的索引,本联通区域任意一点的索引，为了增加可理解性，用最小索引
  int m_iConnetRegionCount;
};
class Solution {
public:
  int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
    m_c = graph.size();
    set<int> setInit(initial.begin(), initial.end());
    CUnionFind uf(m_c);
    for (int i = 0; i < m_c; i++)
    {
      if (setInit.count(i))
      {
        continue;
      }
      for (int j = i + 1; j < m_c; j++)
      {
        if (setInit.count(j))
        {
          continue;
        }
        if (graph[i][j])
        {
          uf.Union(i, j);
        }
      }
    }
    unordered_map<int, int> mRegionInit;
    for (int i = 0; i < m_c; i++)
    {
      if (!setInit.count(i))
      {
        continue;
      }
      for (int j = 0; j < m_c; j++)
      {
        if (setInit.count(j))
        {
          continue;
        }
        if (!graph[i][j])
        {
          continue;
        }
        const int region = uf.GetConnectRegionIndex(j);
        if (mRegionInit.count(region) && (mRegionInit[region] != i ))
        {//新旧初始节点必须不同
          mRegionInit[region] = -1;
        }
        else
        {
          mRegionInit[region] = i;
        }
      }
    }
    map<int, int> mSub;
    auto m = uf.GetNodeOfRegion();
    for (const auto& [region, init] : mRegionInit)
    { 
      if (-1 != init)
      {
        mSub[init] += m[region].size();
      }
    }
    int index = -1;
    int iMax = 0;
    for (const auto& [tmp, cnt] : mSub)
    {
      if (cnt > iMax)
      {
        iMax = cnt;
        index = tmp;
      }
    }
    return (-1 == index) ? *setInit.begin() : index;
  }
  int m_c;
};

template<class T,class T2>
void Assert(const T& t1, const T2& t2)
{
  assert(t1 == t2);
}
template<class T>
void Assert(const vector<T>& v1, const vector<T>& v2)
{
  if (v1.size() != v2.size())
  {
    assert(false);
    return;
  }
  for (int i = 0; i < v1.size(); i++)
  {
    Assert(v1[i], v2[i]);
  }
}
int main()
{
  vector<vector<int>> graph;
  vector<int> initial;
  {
    Solution sln;
    graph = { {1,1,0},{1,1,0},{0,0,1} }, initial = { 0,1 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(0, res);
  }
  {
    Solution sln;
    graph = { {1,1,0},{1,1,1},{0,1,1} }, initial = { 0,1 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(1, res);
  }
  {
    Solution sln;
    graph = { {1,1,0,0},{1,1,1,0},{0,1,1,1},{0,0,1,1} }, initial = { 0,1 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(1, res);
  }
  {
    Solution sln;
    graph = { {1,0,0,0,0,0,0,0,0},{0,1,0,0,0,0,0,0,0},{0,0,1,0,1,0,1,0,0},{0,0,0,1,0,0,0,0,0},{0,0,1,0,1,0,0,0,0},{0,0,0,0,0,1,0,0,0},{0,0,1,0,0,0,1,0,0},{0,0,0,0,0,0,0,1,0},{0,0,0,0,0,0,0,0,1} }, initial = { 6,0,4 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(0, res);
  }
  {
    Solution sln;
    graph = { {1,0,0,0,0,0,0,0,1},{0,1,0,1,0,0,0,0,0},{0,0,1,1,0,1,0,0,0},{0,1,1,1,1,0,1,0,0},{0,0,0,1,1,1,0,0,0},{0,0,1,0,1,1,0,0,0},{0,0,0,1,0,0,1,1,0},{0,0,0,0,0,0,1,1,1},{1,0,0,0,0,0,0,1,1} }, initial = { 3,7 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(3, res);
  }
  {
    Solution sln;
    graph ={ {1,0,0,0,0,1,0},{0,1,1,0,0,0,0},{0,1,1,0,0,0,0},{0,0,0,1,0,0,0},{0,0,0,0,1,0,0},{1,0,0,0,0,1,0},{0,0,0,0,0,0,1} }, initial = { 4 };
    auto res = sln.minMalwareSpread(graph, initial);
    Assert(4, res);
  }
}