[BOJ / solved ac] 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\(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

union-find를 이용해 문제를 풀었다.

 

1. 로봇 번호(각 부품별 root 번호)를 표시하기 위한 robot 배열 생성 (인덱스 => 부품 번호, 값 => root, 각 부품별 초기값은 인덱스)

 

2. 지금까지 알아낸 robot(c)의 부품을 표시하기 위한 depth 배열 생성 (인덱스 => 부품 번호, 값 => 부품 개수, 초기값은 1)

depth[root]에 해당 로봇의 부품 개수를 표시

 

3. 부품별 로봇 번호를 찾는 함수 find_group

3-1. 부품 번호 == 로봇 번호면 부품 번호 return

3-2. 부품 번호 != 로봇 번호면 재귀를 이용해 로봇 번호 찾기

3-3. robot[num]에 로봇 번호 갱신 후, 로봇 번호 return

 

4. 같은 로봇 부품별로 묶는 함수 union_group

3-1. 각 부품별 로봇 번호 찾기

3-2. 두 부품의 로봇 번호가 같으면 함수 종료

3-3. 두 부품의 로봇 번호가 다르면, depth[로봇 번호] 값 비교 후, 더 큰 값에 작은 값을 더해줌

3-4. depth[로봇 번호]이 더 작은 값의 root[부품 번호] 값을 갱신

5. I a b 형태의 지시

union_group 함수 실행

6. Q c 형태의 지시

c의 root를 찾고, depth[c]를 출력

 

 

# 78628 KB / 1788 ms

def find_group(num, robot):
    if robot[num] == num:
        return num
    result = find_group(robot[num], robot)
    robot[num] = result
    return result

def union_group(num1, num2, robot, depth):
    group1, group2 = find_group(int(num1), robot), find_group(int(num2), robot)
    if group1 == group2:
        return
    if depth[group1] >= depth[group2]:
        depth[group1] += depth[group2]
        robot[group2] = group1
    else:
        depth[group2] += depth[group1]
        robot[group1] = group2

def practice_command():
    from sys import stdin
    new_input = stdin.readline
    N = int(new_input())
    limit = int(1e6)+1
    robot, depth = list(range(limit)), [1] * limit

    for _ in range(N):
        command = new_input().split()
        if command[0] == 'I':
            union_group(*command[1:], robot, depth)
        else:
            num = int(command[1])
            group = find_group(num, robot)
            print(depth[group])


practice_command()

---

풀이 2

풀이 1에서 robot, depth를 robot으로 합쳐줌

root일 때는 depth에 해당하는 음수값, root가 아닐 때는 root에 해당하는 양수값

# 48284 KB / 1436 ms

def find_group(num, robot):
    if robot[num] < 0:
        return num
    result = find_group(robot[num], robot)
    robot[num] = result
    return result

def union_group(num1, num2, robot):
    group1, group2 = find_group(int(num1), robot), find_group(int(num2), robot)
    if group1 == group2:
        return
    if robot[group1] <= robot[group2]:
        robot[group1] += robot[group2]
        robot[group2] = group1
    else:
        robot[group2] += robot[group1]
        robot[group1] = group2

def practice_command():
    from sys import stdin
    new_input = stdin.readline
    N = int(new_input())
    robot = [-1] * (int(1e6)+1)

    for _ in range(N):
        command = new_input().split()
        if command[0] == 'I':
            union_group(*command[1:], robot)
        else:
            num = int(command[1])
            group = find_group(num, robot)
            print(-robot[group])


practice_command()

 

---

참고

https://ssungkang.tistory.com/entry/Algorithm-%EC%9C%A0%EB%8B%88%EC%98%A8-%ED%8C%8C%EC%9D%B8%EB%93%9CUnion-Find

[Algorithm] 유니온 파인드(Union - Find)