[Python] 18116번 로봇 조립

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

18116번: 로봇 조립

---

 

24/02/21

 

 

분리 집합 문제로, 집합의 크기를 바로 내보내는 쿼리를 어떻게 구현해야 할 지 집중해야 해결할 수 있는 문제다.

---

 

문제 접근 방식:

 

 

결론부터 말하면, 집합의 크기를 바로 내보내는 리스트를 따로 관리해주면 된다.

 

기존 분리 집합 코드에서는 배열의 인덱스는 해당 노드의 값, 배열의 값은 해당 노드를 자식 노드로 가지는 노드의 값, 즉, 부모 노드의 값을 저장한다.

 

나는 여기에서 집합의 개수를 저장하는 배열을 따로 추가해주었다.

 

이 배열은 유니온을 진행할 때 부모를 기준으로 집합의 개수를 저장하도록 한다.

 

즉, 유니온을 할 때 부모가 서로 다른 경우, 집합의 개수를 서로 합치면 되므로 덧셈을 진행한다.

 

부모가 같은 경우에는 따로 합치는 과정을 진행하지 않는다.

 

추가로 시간을 줄이기 위해 출력할 내용을 모두 리스트에 저장하고 print함수를 한 번만 사용하도록 구현했다.

   

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 18116번 로봇 조립
# 유니온 파인드
import sys
input = sys.stdin.readline

parent = [i for i in range(1_000_000+1)]
count = [1 for _ in range(1_000_000+1)]
def find(node):
    if parent[node] != node:
        parent[node] = find(parent[node])
    return parent[node]
def union(node1, node2):
    p1, p2 = find(node1), find(node2)
    if p1 < p2:
        parent[p2] = p1
        count[p1] += count[p2]
    elif p1 > p2:
        parent[p1] = p2
        count[p2] += count[p1]
def Q(c):
    p = find(c)
    return count[p]

ans_li = []
for _ in range(int(input())):
    I = input().rstrip().split()
    command = I[0]
    if command == 'I':
        a, b = int(I[1]), int(I[2])
        union(a, b)
    else:
        c = int(I[1])
        ans_li.append(Q(c))
print(*ans_li, sep = '\n')