백준 18116 : 로봇 조립 (Python)

본문 링크

import sys
input=sys.stdin.readline

def Find(x):

    if x!=disjoint[x]:
        disjoint[x]=Find(disjoint[x])
    return disjoint[x]

def Union(a,b):

    a=Find(a)
    b=Find(b)

    if a>b:
        disjoint[a]=b
        dp[b]+=dp[a]
    elif a<b:
        disjoint[b]=a
        dp[a]+=dp[b]


N=int(input())

disjoint=[0]*(10**6+1)

for i in range(1,10**6+1):
    disjoint[i]=i

dp=[1]*(10**6+1)

check=set()
for i in range(N):
    L=list(input().split())


    if L[0]=="I":
        A=int(L[1]) ; B=int(L[2])
        Union(A, B)
    else:

        A=Find(int(L[1]))
        print(dp[A])

📌 어떻게 접근할 것인가?

유니온 파인드를 사용했습니다. 다만 문제에서 부품의 번호가 꼭 $N$ 이하이라는 조건이 없기 때문에
배열을 처음 만들때 $10^6+1$ 크기만큼 만들어 줍니다. 그리고 같은 부품의 개수도 빠르게 구하기 위해서 $dp=[1]*(10^6+1)$ 로 만들어 주고 $Union$ 할때마다 부품의 개수를 dp[b]+=dp[a] 와 같은 형식으로 더해줍니다.

이때 주의해야할 점은 a==b 일떄는 더해주면 안됩니다.