[Python] 백준 20055 컨베이어 벨트 위의 로봇

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

컨베이어 벨트를 이용해 로봇들을 건너편으로 옮기려고 한다. 로봇을 옮기는 과정에서는 아래와 같은 일이 순서대로 일어난다.

1. 벨트가 각 칸 위에 있는 로봇과 함께 한 칸 회전한다.
2. 가장 먼저 벨트에 올라간 로봇부터, 벨트가 회전하는 방향으로 한 칸 이동할 수 있다면 이동한다. 만약 이동할 수 없다면 가만히 있는다.
   1. 로봇이 이동하기 위해서는 이동하려는 칸에 로봇이 없으며, 그 칸의 내구도가 1 이상 남아 있어야 한다.
3. 올리는 위치에 있는 칸의 내구도가 0이 아니면 올리는 위치에 로봇을 올린다.
4. 내구도가 0인 칸의 개수가 K개 이상이라면 과정을 종료한다. 그렇지 않다면 1번으로 돌아간다.

종료되었을 때 몇 번째 단계가 진행 중이었는지 구해보자. 가장 처음 수행되는 단계는 1번째 단계이다.

입력
첫째 줄에 N, K가 주어진다. 둘째 줄에는 A1, A2, ..., A2N이 주어진다.

출력
몇 번째 단계가 진행 중일때 종료되었는지 출력한다.

제한
2 ≤ N ≤ 100
1 ≤ K ≤ 2N
1 ≤ Ai ≤ 1,000

내 첫 번째 제출은 다음과 같다.

python3에선 시간 초과가 발생하였고, pypy에선 통과되었다.

# 20055 컨베이어 벨트 위의 로봇

def check_N():
    if belts_box[N - 1] != 0:
        belts_box[N - 1] = 0
    return


N, K = map(int, input().split())
belts_strength = list(map(int, input().split()))
belts_box = [0] * (2 * N)
length = len(belts_strength)
cnt = 1

while True:
    check_N()
    # 1. 회전
    for i in range(length - 1, 0, -1):
        temp = belts_strength[i]
        belts_strength[i] = belts_strength[i - 1]
        belts_strength[i - 1] = temp

        temp = belts_box[i]
        belts_box[i] = belts_box[i - 1]
        belts_box[i - 1] = temp
    
    check_N()
    # 2. 이동
    for i in range(N - 2, -1, -1):
        if belts_box[i] > 0 and belts_box[i + 1] == 0 and belts_strength[i + 1] > 0:
            belts_box[i + 1] = belts_box[i]
            belts_box[i] = 0
            belts_strength[i + 1] -= 1

    check_N()
    # 3. 로봇 올리기
    if belts_strength[0] != 0:
        belts_strength[0] -= 1
        belts_box[0] += 1
    
    check_N()
    # 4. 내구도 0인 칸의 갯수 K 이상이면 종료.
    if belts_strength.count(0) >= K:
        break
    cnt += 1
        
print(cnt)

맨 처음엔 2. 이동 부분의 for문에서 0번 째 벨트부터 N-1번 째 벨트의 순서로 탐색을 진행했으나, 예시 코드에서 틀린 방식임을 깨닫고 N-2번 째부터 0번까지 탐색하도록 알고리즘을 변경하였다.

다른 분들의 코드를 참고해보니, 1. 회전 부분에서 for문을 사용하지 않아도 된다는 것을 알게 되었다.

check_N 함수도 1. 회전, 2. 이동 이후에만 검사해도 괜찮았다.

# 20055 컨베이어 벨트 위의 로봇

def check_N():
    if belts_box[N - 1] != 0:
        belts_box[N - 1] = 0
    return


N, K = map(int, input().split())
belts_strength = list(map(int, input().split()))
belts_box = [0] * (2 * N)
length = len(belts_strength)
cnt = 1
temp_list = list()

while True:
    # 1. 회전
    temp_list = [belts_strength[-1]] + belts_strength[:-1]
    belts_strength = temp_list[:]
    
    temp_list = [belts_box[-1]] + belts_box[:-1]
    belts_box = temp_list[:]
    
    check_N()
    # 2. 이동
    for i in range(N - 2, -1, -1):
        if belts_box[i] > 0 and belts_box[i + 1] == 0 and belts_strength[i + 1] > 0:
            belts_box[i + 1] = belts_box[i]
            belts_box[i] = 0
            belts_strength[i + 1] -= 1

    check_N()
    # 3. 로봇 올리기
    if belts_strength[0] != 0:
        belts_strength[0] -= 1
        belts_box[0] += 1
    
    # 4. 내구도 0인 칸의 갯수 K 이상이면 종료.
    if belts_strength.count(0) >= K:
        break
    cnt += 1
        
print(cnt)

결론 및 느낀점

• 흥미로운 구현 문제였다. 맨 처음 문제를 보았을 때, 해석하느라 애먹었다.
• 단순 for문이 아닌, slicing을 통해서 벨트의 회전을 구현할 수 있었다.