문제

18116번: 로봇 조립

풀이

I a b 형태로 들어오는 a, b 부품은 같은 로봇이 부품이고, 서로 다른 로봇은 공동 부품을 가지지 않기 때문에 Union Find 알고리즘으로 해결할 수 있다.
두 부품 a, b를 한 집합으로 처리하며 Q c의 형태로 들어오게 되면 c에 해당하는 집합의 원소 개수를 출력하면 된다.

원소의 개수를 매번 세는 것은 비효율적이기 때문에 미리 개수를 계산하여 저장해둬야 한다.
Union 연산을 하며 두 집합의 원소 개수를 합친 후 그 값을 루트 노드 위치에 저장해두면 매번 셀 필요 없이 해당 부품의 루트에서 원소 개수를 알 수 있다.

코드

#include <iostream>
#include <vector>

#define MAX 1e6 + 1
using namespace std;

int n;
vector<int> root(MAX, 0);
vector<int> depth(MAX, 0);
vector<int> partsCount(MAX, 1);

int Find(int x)
{
    if (root[x] == x)
    {
        return x;
    }
    else
    {
        return root[x] = Find(root[x]);
    }
}

void Union(int x, int y)
{
    x = Find(x);
    y = Find(y);

    if (x==y) return;

    if (depth[x] > depth[y])
    {
        root[y] = x;
        partsCount[x] += partsCount[y];
    }
    else if (depth[y] > depth[x])
    {
        root[x] = y;
        partsCount[y] += partsCount[x];
    }
    else{
        root[x] = y;
        depth[y]++;
        partsCount[y] += partsCount[x];
    }
}

int main ()
{
    cin.tie(0);
    cout.tie(0);
    ios_base::sync_with_stdio(false);
    for (uint32_t i=0; i<root.size(); i++)
    {
        root[i] = i;
    }
    cin >> n;
    while (n--)
    {
        char command;
        cin >> command;
        if (command == 'I')
        {
            int x, y;
            cin >> x >> y;
            Union(x, y);
        }
        else
        {
            int x;
            cin >> x;
            cout << partsCount[Find(x)] << "\n";
        }
    }
    return 0;
}