[Python] 백준 / gold / 18116 : 로봇 조립

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

Disjoint Set (분리 집합) 문제. 웰노운이듯이 Union-Find 메서드를 작성해서 해결하면 된다. 근데 이 문제는 집합의 크기를 묻는 쿼리도 있기 때문에 set_size 라는 크기 저장 리스트를 따로 선언해줬다.

이 문제를 푸는 데 주의할 게 2가지 정도 있었던 거 같다.
1. set_size 를 계산하기 위해 root node 가 같은 경우는 어떻게 처리할 것인지 고민할 것
2. Python3 이 아닌 Pypy3 으로 제출할 것..

파이썬 코드

import sys
N = int(sys.stdin.readline())
cmd = []

for _ in range(N):
    cmd.append(sys.stdin.readline().split())

parent = [x for x in range(10**6 + 1)]
set_size = [1 for x in range(10**6 + 1)]


def find_root(x):
    if parent[x] == x:
        return x
    else:
        parent[x] = find_root(parent[x])
        return parent[x]


def union(a, b):
    pa = find_root(a)
    pb = find_root(b)

    if pa < pb:
        parent[pb] = pa
        set_size[pa] += set_size[pb]
        set_size[pb] = 0
    elif pb < pa:
        parent[pa] = pb
        set_size[pb] += set_size[pa]
        set_size[pa] = 0
    # 서로 이미 루트 노드가 같다면
    else:
        pass


for query in cmd:
    # 합집합
    if query[0] == 'I':
       union(int(query[1]), int(query[2]))
    # 질문
    elif query[0] == 'Q':
        c = int(query[1])
        pc = find_root(c)
        print(set_size[pc])