[BOJ] [Python] 18116 : 로봇 조립

Union-Find 알고리즘으로 해결했다.

 

이 문제의 핵심은 루트 노드 밑에 몇 개의 노드가 존재하는지 파악하는 것이다.

 

 

import sys

input = sys.stdin.readline

n = int(input())

parent = [-1 for i in range(1000000 + 1)]


def find(x):
    if parent[x] < 0:
        return x

    else:
        y = find(parent[x])
        parent[x] = y

        return y


def union(x, y):
    x = find(x)
    y = find(y)

    if x == y:
        return

    if parent[x] < parent[y]:
        A, B = x, y

    else:
        A, B = y, x

    parent[A] += parent[B]
    parent[B] = A


for _ in range(n):
    li = list(map(str, input().split()))

    if li[0] == "I":
        a = int(li[1])
        b = int(li[2])

        union(a, b)

    if li[0] == "Q":
        a = int(li[1])

        print(parent[find(a)] * - 1)

 

 

예제 입력 1

 

예제 입력 1로 설명하겠다.

 

x, y = 1, 2

x, y = 3, 2

 

따로 조건이 없으므로

x가 속한 트리와 y가 속한 트리, 둘 중 어느 것이 밑으로 가게 되던 상관없다.

하지만 효율성을 생각해서, 트리의 크기가 작은 것이 트리의 크기가 큰 것 밑으로 가게 했다.

 

만약 트리의 크기가 같다면, y가 속한 트리 밑으로 x가 속한 트리가 들어가게 했다. (이유 없음)

 

parent[2] = -3

parent[1] = 2

parent[3] = 2

 

노드 2는 현재 루트 노드이다. 따라서 트리의 크기에 대한 정보를 저장한다.

하위 노드는 루트 노드를 가리키고 있다.