Advertisement
den4ik2003

Untitled

Sep 19th, 2023
728
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.81 KB | None | 0 0
  1. #include <iostream>
  2. #include <vector>
  3. #include <algorithm>
  4.  
  5. using std::vector;
  6.  
  7. class DSU {
  8.  public:
  9.   explicit DSU(int n) : parent_(vector<int>(n, -1)),
  10.                         size_(vector<int>(n, 0)),
  11.                         whenCame_(vector<int>(n, 1)),
  12.                         numberOfUnits_(vector<int>(n, 1)) {} // TODO - точно ли 1 тут
  13.  
  14.   int Get(int v) {
  15.     if (parent_[v] == -1) {
  16.       return v;
  17.     }
  18.     return Get(parent_[v]);
  19.   }
  20.  
  21.   bool CheckTwoElements(int u, int v) {
  22.     return Get(v) == Get(u);
  23.   }
  24.  
  25.   void Unit(int u, int v) {
  26.     u = Get(u);
  27.     v = Get(v);
  28.     if (u != v) {
  29.       if (size_[u] < size_[v]) { // v родитель u
  30.         parent_[u] = v;
  31.         size_[v] += size_[u];
  32.         whenCame_[u] = numberOfUnits_[v];
  33.         ++numberOfUnits_[v];
  34.       } else { // u родитель v
  35.         parent_[v] = u;
  36.         size_[u] += size_[v];
  37.         whenCame_[v] = numberOfUnits_[u];
  38.         ++numberOfUnits_[u];
  39.       }
  40.     }
  41.   }
  42.  
  43.   int GetNumberOfUnits(int u) {
  44.     if (parent_[u] == -1) {
  45.       return numberOfUnits_[u];
  46.     }
  47.     int uParent = Get(parent_[u]);
  48.     return numberOfUnits_[uParent] - whenCame_[u] + numberOfUnits_[u];
  49.   }
  50.  
  51.  private:
  52.   vector<int> parent_;
  53.   vector<int> size_;
  54.   vector<int> whenCame_;
  55.   vector<int> numberOfUnits_;
  56. };
  57.  
  58. int main() {
  59.   int n, m;
  60.   std::cin >> n >> m;
  61.   DSU dsu(n);
  62.   for (int i = 0; i < m; ++i) {
  63.     int type, x;
  64.     std::cin >> type >> x;
  65.     if (type == 1) {
  66.       int y;
  67.       std::cin >> y;
  68.       dsu.Unit(x - 1, y - 1);
  69.     } else if (type == 2) {
  70.       int y;
  71.       std::cin >> y;
  72.       if (dsu.CheckTwoElements(x - 1, y - 1)) std::cout << "YES\n";
  73.       else std::cout << "NO\n";
  74.     } else if (type == 3) {
  75.       std::cout << dsu.GetNumberOfUnits(x - 1) << "\n";
  76.     }
  77.   }
  78. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement