문제

초기에 {0}, {1}, {2}, … {n} 이 각각 n+1개의 집합을 이루고 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

입력

첫째 줄에 n(1≤n≤1,000,000), m(1≤m≤100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

출력

  • 1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

예제

예제 입력
7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

예제 출력
NO
NO
YES

접근

  • Disjoint-Set(Union-Find)를 공부하기 위해 찾은 문제다 보니 자연스럽게 이 방법으로 접근하게 되었다.

내가 생각한 풀이

  • 합집합 연산이 이뤄질 경우 merge 과정을 통해 합쳐주고 확인하는 과정일 때는 find를 통해 찾은 root 값이 동일한지 비교한다.

미처 생각하지 못했던것

  • 문제의 입력값이 굉장히 많다는 것을 고려하여 입력 속도를 빠르게 해줬어야 했는데 이를 고려하지 않아 계속 시간초과가 났었다.

코드

#include<iostream>
#include<vector>
using namespace std;
struct OptimizedDisjointSet {
    vector<int> parent, rank;
    OptimizedDisjointSet (int n) : parent(n), rank(n,1) {
        for (int i=0 ; i < n; i++)
        {
          parent[i] = i;
        }
    }

    int find(int u) {
        if (u == parent[u]) return u;
        return parent[u] = find(parent[u]); //find 최적화를 위한 작업, 부모가 같은 것들에 대해 동일한 부모로 바꿔준다.
    }

    void merge (int u, int v) {
        u = find(u); v = find(v);
        if (u == v) return;
        if (rank[u] > rank[v]) swap(u, v); //tree의 높이를 줄이기 위해 rank를 설정한다.
        parent[u] = v;
        if (rank[u] == rank[v]) ++rank[v];
    }
};
int main(void)
{
  //입력값이 매우 많기 때문에 입력 속도에서 시간을 줄여줘야 시간초과가 나지 않는다.
  ios_base :: sync_with_stdio(false);
  cin.tie(NULL);
  int n, m;
  cin>>n>>m;
  struct OptimizedDisjointSet DJS(n+1);
  for(int i = 0 ; i<m ; i++)
  {
    int t, a, b;
    cin>>t>>a>>b;
    if(t == 0)
    {
      DJS.merge(a, b);
    }
    else
    {
      int t1, t2;
      t1 = DJS.find(a);
      t2 = DJS.find(b);
      if(t1 == t2) cout<<"YES\n";
      else cout<<"NO\n";
    }
  }
  return 0;
}