[백준 / BOJ] C++ 18116 로봇 조립

18116번: 로봇 조립

 

문제

https://www.acmicpc.net/problem/18116

18116번: 로봇 조립

 

 

풀이

같은 로봇의 부품끼리 집합을 만들어 집합의 크기를 구할 수 있어야 하므로 유니온 파인드를 사용한 분리 집합으로 풀 수 있다. 부품이 1부터 10^6까지 표현되므로 배열의 크기를 10^6으로 고정시켜야 한다. 처음에 각 부품은 자기 자신 혼자의 집합이므로 크기가 1이다. 따라서 집합의 원소의 수를 나타내는 cnt 배열을 1로 초기화해준다.

 

I 쿼리가 입력되었을 때는 a와 b를 같은 집합으로 묶어야 하므로 각각의 루트 노드를 구해 크기를 한쪽으로 합쳐주는데 항상 작은 번호 쪽으로 합쳐준다. 그리고 큰 번호의 루트 노드를 작은 번호의 루트 노드로 설정해 같은 집합으로 만들어 준다. 단, 같은 집합에 속한 경우에는 이 작업을 하지 않는다.

 

Q 쿼리가 입력되었을 때는 c가 포함된 집합의 크기를 물어보는 것이므로 c의 루트 노드 인덱스인 cnt 배열의 원소를 출력하면 된다.

 

 

코드

#include <bits/stdc++.h>
using namespace std;

vector<int> p(1000001), cnt(1000001, 1);

int find(int x) {
  if (p[x] != x)
    p[x] = find(p[x]);
  return p[x];
}

void merge(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y)
    return;
  if (x > y)
    swap(x, y);
  cnt[x] += cnt[y];
  cnt[y] = 0;
  p[y] = x;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  int n;
  cin >> n;
  iota(p.begin(), p.end(), 0);
  while (n--) {
    char c;
    int a, b;
    cin >> c;
    if (c == 'I') {
      cin >> a >> b;
      merge(a, b);
    } else {
      cin >> a;
      cout << cnt[find(a)] << '\n';
    }
  }
  return 0;
}