[백준 18116번] 로봇 조립 (C++)

문제링크 : https://www.acmicpc.net/problem/18116

18116번: 로봇 조립

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int MAX = 987654321;

int N, M;
int parent[1000001];
int cnt[1000001];

int find_root(int x) //부모찾기
{
    if(x == parent[x]) return x;
    return parent[x] = find_root(parent[x]);
}

void Union(int a, int b) //병합
{
    a=find_root(a);
    b=find_root(b);
    if(a!=b)
    {
        cnt[a]+=cnt[b]; //부모에 값을 연결함과 동시에 부모에 자식값 더해주기
        parent[b]=a;
    }
}


int main()
{
    ios_base::sync_with_stdio(0); 
	cin.tie(0);
    
    cin >> N;

    for(int i=1; i<=1000000; i++) //초기화
    {
        parent[i]=i;
        cnt[i]=1;
    }

    while(N--)
    {
        char c; int a, b;
        cin >> c;
        if(c=='I') 
        {
            cin >> a >> b;
            Union(a, b);
        }
        else 
        {
            cin >> a;
            cout << cnt[find_root(a)] << "\n"; //부모에 저장된 갯수 출력
        }
    }
}

 

살짝 응용된 분리집합 문제이다.

분리집합 알고리즘에서 병합과 부모 찾기만 필요하고, 이어졌는지 여부는 확인할 필요가 없으므로 해당 함수는 제거했다.

 

먼저 parent 배열을 초기화 할 때 갯수를 담을 cnt 배열도 같이 각 1로 초기화를 해준다.

이후 병합과정에서 값을 연결함과 동시에 부모쪽에 자식의 갯수를 더해준다.

이후 출력할때 출력할 부모의 총 갯수를 출력해주면 된다.

4
I 1 2
I 3 2
Q 1

cnt[1] += cnt[2] -> 2
cnt[3] += cnt[1(2의 부모)] -> 3

find_root(1) = 3
cnt[3] = 3

 

병합시 주의할 점은 병합 대상인 a와 b가 다를 때만 더해주는 작업을 해주어야 한다.

그렇지 않으면 이미 더해진 값이 또 더해지는 중복 상황이 발생한다.

6
I 1 2
I 1 3
I 2 3
Q 1
Q 2
Q 3

옳은 답 : 3 3 3
중복 발생 : 6 6 6

 

처음에 중복을 생각안하여 여러번 틀렸던 문제이다.