백준 18116 [로봇 조립] 파이썬

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

 

# 백준 18116 로봇 조립
# Union-Find (Disjoint Set Union, DSU) + Union by Size 최적화

import sys
sys.setrecursionlimit(10**7)   # 재귀 깊이 제한 증가 (find 함수에서 재귀 사용 가능)
input = sys.stdin.readline
N = int(input())  # 명령 개수


# find 연산 (경로 압축 포함)
# x의 루트(대표 노드)를 찾음
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])  # 경로 압축
    return parent[x]


# union 연산 (union by size)
# 두 집합을 합칠 때, 더 작은 집합을 큰 집합에 붙인다.
def union(a, b):
    a = find(a)  # a의 루트
    b = find(b)  # b의 루트
    if a != b:   # 루트가 다르면 합치기
        if sizes[a] < sizes[b]:
            parent[a] = b
            sizes[b] += sizes[a]  # b 집합 크기에 a 집합 크기 추가
        else:
            parent[b] = a
            sizes[a] += sizes[b]  # a 집합 크기에 b 집합 크기 추가


# parent: 각 원소의 부모 (처음엔 자기 자신이 부모)
# sizes: 각 루트가 대표하는 집합의 크기
sizes = [1] * (10**6 + 1)
parent = [i for i in range(10**6 + 1)]


# 명령 처리
for _ in range(N):
    order = input().split()  # "I a b" 또는 "Q x"
    if order[0] == 'I':  
        # I a b → 로봇 a와 b를 연결 (union)
        union(int(order[1]), int(order[2]))
    else:
        # Q x → x가 속한 집합의 크기 출력
        print(sizes[find(int(order[1]))])

UNION-FIND를 이용하되, UNION by size, 즉 size를 기준으로 UNION을 진행하였다.