【树上倍增】【割点】 【换根法】3067. 在带权树网络中统计可连接服务器对数目（三）

class CNeiBo
{
public: 
  static vector<vector<int>> Two(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0) 
  {
    vector<vector<int>>  vNeiBo(n);
    for (const auto& v : edges)
    {
      vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase);
      if (!bDirect)
      {
        vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase);
      }
    }
    return vNeiBo;
  } 
  static vector<vector<std::pair<int, int>>> Three(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0)
  {
    vector<vector<std::pair<int, int>>> vNeiBo(n);
    for (const auto& v : edges)
    {
      vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase, v[2]);
      if (!bDirect)
      {
        vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase, v[2]);
      }
    }
    return vNeiBo;
  }
};
class CUnionFind
{
public:
  CUnionFind(int iSize) :m_vNodeToRegion(iSize)
  {
    for (int i = 0; i < iSize; i++)
    {
      m_vNodeToRegion[i] = i;
    }
    m_iConnetRegionCount = iSize;
  } 
  CUnionFind(vector<vector<int>>& vNeiBo):CUnionFind(vNeiBo.size())
  {
    for (int i = 0; i < vNeiBo.size(); i++) {
      for (const auto& n : vNeiBo[i]) {
        Union(i, n);
      }
    }
  }
  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 CParents
{
public:
  CParents(vector<int>& vParent, const int iMaxLeve)
  { 
    int iBitNum = 0;
    for (; (1 << iBitNum) < iMaxLeve; iBitNum++);
    const int n = vParent.size();
    m_vParents.assign(iBitNum+1, vector<int>(n, -1));
    m_vParents[0] = vParent;
    //树上倍增
    for (int i = 1; i < m_vParents.size(); i++)
    {
      for (int j = 0; j < n; j++)
      {
        const int iPre = m_vParents[i - 1][j];
        if (-1 != iPre)
        {
          m_vParents[i][j] = m_vParents[i - 1][iPre];
        }
      }
    }
  }
  int GetParent(int iNode, int iLeve)const
  {
    int iParent = iNode;
    for (int iBit = 0; iBit < m_vParents.size(); iBit++)
    {
      if (-1 == iParent)
      {
        return iParent;
      }
      if (iLeve & (1 << iBit))
      {
        iParent = m_vParents[iBit][iParent];
      }
    }
    return iParent;
  } 
protected:
  vector<vector<int>> m_vParents;
};
class C2Parents : CParents
{
public:
  C2Parents(vector<int>& vParent, const vector<int>& vLeve) :m_vLeve(vLeve)
    , CParents(vParent,*std::max_element(vLeve.begin(), vLeve.end()))
  {   
  } 
  int GetPublicParent(int iNode1, int iNode2)const
  {
    int leve0 = m_vLeve[iNode1];
    int leve1 = m_vLeve[iNode2];
    if (leve0 < leve1)
    {
      iNode2 = GetParent(iNode2, leve1 - leve0);
      leve1 = leve0;
    }
    else
    {
      iNode1 = GetParent(iNode1, leve0 - leve1);
      leve0 = leve1;
    }
    //二分查找
    int left = -1, r = leve0;
    while (r - left > 1)
    {
      const auto mid = left + (r - left) / 2;
      const int iParent0 = GetParent(iNode1, mid);
      const int iParent1 = GetParent(iNode2, mid);
      if (iParent0 == iParent1)
      {
        r = mid;
      }
      else
      {
        left = mid;
      }
    }
    return GetParent(iNode1, r);
  }
protected:
  vector<vector<int>> m_vParents;
  const vector<int> m_vLeve;
};
//割点
class CCutPoint
{
public:
  CCutPoint(const vector<vector<int>>& vNeiB) : m_iSize(vNeiB.size())
  {
    m_vNodeToTime.assign(m_iSize, -1);
    m_vCutNewRegion.resize(m_iSize);
    for (int i = 0; i < m_iSize; i++)
    {
      if (-1 == m_vNodeToTime[i])
      {
        m_vRegionFirstTime.emplace_back(m_iTime);
        dfs(vNeiB, i, -1);
      }
    } 
  }
  int dfs(const vector<vector<int>>& vNeiB,const int cur, const int parent)
  {
    int iMinTime = m_vNodeToTime[cur] = m_iTime++;
    int iRegionCount = (-1 != parent);//根连通区域数量
    for (const auto& next : vNeiB[cur])   {
      if (-1  != m_vNodeToTime[next])     {
        iMinTime = min(iMinTime, m_vNodeToTime[next]);
        continue;
      }
      const int childMinTime = dfs(vNeiB, next, cur);
      iMinTime = min(iMinTime, childMinTime);
      if (childMinTime >= m_vNodeToTime[cur])     {
        iRegionCount++;
        m_vCutNewRegion[cur].emplace_back(m_vNodeToTime[next], m_iTime);
      }
    }
    if (iRegionCount < 2)
    {
      m_vCutNewRegion[cur].clear();
    }
    return iMinTime;
  }
  const int m_iSize;
  const vector<int>& Time()const { return m_vNodeToTime; }//各节点的时间戳
  const vector<int>& RegionFirstTime()const { return m_vRegionFirstTime; }//各连通区域的最小时间戳
  vector<bool> Cut()const { 
    vector<bool> ret;
    for (int i = 0; i < m_iSize; i++)
    {
      ret.emplace_back(m_vCutNewRegion[i].size());
    }
    return ret; }//
  const vector < vector<pair<int, int>>>& NewRegion()const { return m_vCutNewRegion; };
protected:
  vector<int> m_vNodeToTime;
  vector<int> m_vRegionFirstTime;
  vector < vector<pair<int, int>>> m_vCutNewRegion; //m_vCutNewRegion[c]如果存在[left,r) 表示割掉c后，时间戳[left,r)的节点会形成新区域
  int m_iTime = 0;
};
class CConnectAfterCutPoint 
{
public:
  CConnectAfterCutPoint(const vector<vector<int>>& vNeiB) :m_ct(vNeiB)
  {
    m_vTimeToNode.resize(m_ct.m_iSize);
    m_vNodeToRegion.resize(m_ct.m_iSize);
    for (int iNode = 0; iNode < m_ct.m_iSize; iNode++)
    {
      m_vTimeToNode[m_ct.Time()[iNode]] = iNode;
    }
    for (int iTime = 0,iRegion= 0; iTime < m_ct.m_iSize; iTime++)
    {
      if ((iRegion < m_ct.RegionFirstTime().size()) && (m_ct.RegionFirstTime()[iRegion] == iTime))
      {
        iRegion++;
      }
      m_vNodeToRegion[m_vTimeToNode[iTime]] = (iRegion - 1);
    }
  }
  bool Connect(int src, int dest, int iCut)const
  {
    if (m_vNodeToRegion[src] != m_vNodeToRegion[dest])
    {
      return false;//不在一个连通区域
    }
    if (0 == m_ct.NewRegion()[iCut].size())
    {//不是割点
      return true;
    }
    const int r1 = GetCutRegion(iCut, src);
    const int r2 = GetCutRegion(iCut, dest);
    return r1 == r2;
  }
  vector<vector<int>> GetSubRegionOfCut(const int iCut)const
  {//删除iCut及和它相连的边后，iCut所在的区域会分成几个区域：父节点一个区域、各子节点   一个区域
      //父节点所在区域可能为空，如果iCut所在的连通区域只有一个节点，则返回一个没有节点的      区域。
    const auto& v = m_ct.NewRegion()[iCut];
    vector<int> vParen;
    const int iRegion = m_vNodeToRegion[iCut];
    const int iEndTime = (iRegion + 1 == m_ct.RegionFirstTime().size()) ? m_ct.m_iSize : m_ct.RegionFirstTime()[iRegion+1];
    vector<vector<int>> vRet; 
    for (int iTime = m_ct.RegionFirstTime()[iRegion],j=-1; iTime < iEndTime; iTime++)
    {
      if (iCut == m_vTimeToNode[iTime])
      {
        continue;
      }
      if ((j + 1 < v.size()) && (v[j + 1].first == iTime))
      {
        j++;
        vRet.emplace_back();
      }
      const int iNode = m_vTimeToNode[iTime];
      if ((-1 != j ) && (iTime >= v[j].first) && (iTime < v[j].second))
      {
        vRet.back().emplace_back(iNode);
      }
      else
      {
        vParen.emplace_back(iNode);
      }     
    }
    vRet.emplace_back();
    vRet.back().swap(vParen);
    return vRet;
  } 
protected:
  int GetCutRegion(int iCut, int iNode)const
  {
    const auto& v = m_ct.NewRegion()[iCut];
    auto it = std::upper_bound(v.begin(), v.end(), m_ct.Time()[iNode], [](int time, const std::pair<int, int>& pr) {return  time < pr.first; });
    if (v.begin() == it)
    {
      return v.size();
    }
    --it;
    return (it->second > m_ct.Time()[iNode]) ? (it - v.begin()) : v.size();
  }
  vector<int> m_vTimeToNode;
  vector<int> m_vNodeToRegion;//各节点所在区域
  const CCutPoint m_ct;
};
class Solution {
public:
  vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
    m_c = edges.size() + 1;
    m_vDisToRoot.resize(m_c);
    m_vParent.resize(m_c);
    m_vLeve.resize(m_c);
    auto neiBo = CNeiBo::Three(m_c, edges, false, 0);
    DFS(neiBo, 0, -1, 0, 0);
    C2Parents par(m_vParent, m_vLeve);
    auto neiBo2 = CNeiBo::Two(m_c, edges, false, 0);
    CConnectAfterCutPoint cut(neiBo2);
    vector<int> vRet(m_c);
    for (int c = 0; c < m_c; c++)
    {
      auto regs = cut.GetSubRegionOfCut(c);
      int left = 0;
      int& iRet = vRet[c];
      for (const auto& subRegion : regs)
      {
        int cur = 0;
        for (const auto& ab : subRegion)
        {
          const int pub = par.GetPublicParent(ab, c);
          const int len = m_vDisToRoot[ab] + m_vDisToRoot[c] - 2 * m_vDisToRoot[pub];
          if (0 != len % signalSpeed)
          {
            continue;
          }
          cur++;
        }
        iRet += left * cur;
        left += cur;
      }
    }
    return vRet;
  }
  void DFS(vector<vector<std::pair<int, int>>>& neiBo, int cur, int par, int leve, int dis)
  {
    m_vDisToRoot[cur] = dis;
    m_vParent[cur] = par;
    m_vLeve[cur] = leve;
    for (const auto& [next, len] : neiBo[cur])
    {
      if (next == par)
      {
        continue;
      }
      DFS(neiBo, next, cur, leve + 1, dis + len);
    }
  }
  vector<int> m_vDisToRoot, m_vParent, m_vLeve;
  int m_c;
};