[BOJ/Python] 18116_로봇 조립

문제 출처:https://www.acmicpc.net/problem/18116

18116번: 로봇 조립

문제

성규는 로봇을 조립해야 한다. 상자 안에는 여러 로봇의 부품들이 섞여 있다. 그런데 어떤 부품이 어느 로봇의 부품인지 표시가 되어있지 않다. 호재는 전자과라서 두 부품을 보면 같은 로봇의 부품인지 알 수 있다. 그래서 성규는 호재의 지시에 따라 부품들을 정리하기로 하였다.

부품들은 1부터 10^6까지의 정수로 표현된다. 그리고 부품 i가 속한 로봇은 robot(i)라고도 표현한다. 예를 들어, 부품 11과 부품 22가 로봇 A의 부품이라고 알고 있는 경우, robot(11)은 로봇 A를 의미하고, robot(22)도 로봇 A를 의미한다.

서로 다른 로봇은 공통 부품을 가지지 않는다. 즉 어떤 부품이 로봇 A의 부품이라면, 로봇 B의 부품은 될 수 없다.

호재는 2가지 지시를 한다.

• 서로 다른 부품 2개를 말해주며, 두 부품은 같은 로봇의 부품이라는 정보를 알려준다.
• 부품 i에 대해서, 지금까지 알아낸 robot(i)의 부품이 몇 개냐고 물어본다.

초기에는 부품에 대한 정보가 존재하지 않는다.

입력

첫 번째 줄에 호재의 지시 횟수 N이 들어온다. (1 ≤ N ≤ 106)

다음 줄부터 N개의 지시가 들어온다.

부품 2개가 같은 로봇의 부품인지 알려줄 때에는 I a b 의 형태로 들어온다. 부품 a와 부품 b는 같은 로봇의 부품이라는 의미이다. (1 ≤ a, b ≤ 10^6, a ≠ b, a, b는 정수)

어떤 로봇의 부품이 몇 개인지 물어볼 때에는 Q c 의 형태로 들어온다. 지금까지 알아낸 robot(c)의 부품이 몇 개냐는 의미이다. (1 ≤ c ≤ 10^6, c는 정수)

입력으로 Q c의 형태가 적어도 한 번 들어온다.

출력

Q로 시작하는 입력에 대해서 한 줄에 하나씩, 지금까지 알아낸 해당 로봇의 부품 개수를 출력한다.

예제 입력 1

4
I 1 2
I 3 2
Q 1
Q 4

예제 출력 1

3
1

---

풀이

 

1. 부품의 종류는 1부터 1,000,000까지 존재한다.
2. I a b 의 입력의 경우, a 부품과 b 부품은 같은 로봇의 부품임을 의미한다.
3. Q a 의 입력의 경우, a 부품을 이용하여 만드는 로봇에 대하여 현재까지 알아낸 해당 로봇의 부품의 개수를 출력함을 의미한다.

풀이는 다음과 같다.

 

1. 우선 부품의 종류가 1,000,000 까지 이므로, 1,000,000 + 1 까지의 리스트를 두 개 생성한다.
2. 하나는 부품들이 같은 로봇의 부품임을 처리하기 위하여 union & find를 사용하는 데 쓰는 parent 리스트이다.
3. 다른 하나는 Q a의 입력이 들어왔을 때, 현재까지 알아낸 해당 로봇의 부품의 개수를 저장할 robot 리스트이다.
4. I a b 입력이 들어왔을 때, a 와 b를 union_parent 함수를 이용하여 하나의 집합으로 합친다.
5. 함수 내부에서 find_parent함수를 이용하여 a와 b 의 부모 중 루트인 a`과 b`을 찾는다.
6. 만약 a`이 b` 의 부모가 된다면, robot[a`] 에 robot[b`] 값을 더해준다.
7. Q a의 입력이 들어오게 되면, find_parent 함수를 이용하여 a의 부모인 a`을 찾고, robot[a`]을 출력한다.

 

코드는 다음과 같다.

 

import sys
 
def find_parent(parent,x):
    if parent[x] != x:
        parent[x] = find_parent(parent,parent[x])
    return parent[x]
 
def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if a<b:
        parent[b] = a
        robot[a] += robot[b]
    elif a>b:
        parent[a] = b
        robot[b] += robot[a]
 
 
n = int(sys.stdin.readline())
 
parent = [i for i in range(1000000+1)]
robot = [1 for _ in range(1000000+1)]
 
for _ in range(n):
    l = list(sys.stdin.readline().split())
    if l[0]=='I':
        union_parent(parent,int(l[1]),int(l[2]))
    else:
        print(robot[find_parent(parent,int(l[1]))])