백준 컨베이어 벨트 위의 로봇 코드 및 해설 (파이썬)

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

20055번: 컨베이어 벨트 위의 로봇

 

각 칸의 내구도를 belt, 로봇이 있는 칸의 인덱스를 robots에 저장한다. 

1-4 각 과정을 함수로 만들어 매 단계마다 반복한다.

먼저 컨베이어 벨트의 회전은 벨트와 그 위의 로봇이 이동해야 한다. 벨트의 회전은 큐를 이용해 구현했고, 로봇의 이동은 인덱스를 하나씩 증가시키는 것으로 구현했다. 이때 마지막 인덱스에 로봇이 있는 경우 +1이 아니라 처음 칸인 0으로 가야하므로 +1을 하고 컨베이어 벨트의 총 길이인 2N으로 나눈 나머지를 취해 이동할 인덱스를 구했다.

이때 이동할 칸이 내리는 칸(N-1)이면 바로 내려야 하므로, 이 값은 버린다. 

두 번째 과정인 로봇의 이동은 먼저 컨베이어 벨트에 올라온 로봇부터 처리해야 한다. 로봇을 리스트로 만들고 새로 올라온 로봇은 뒷부분에 append 하는 식으로 구현했으므로 robots 리스트 순서대로 하나씩 처리하면 된다. 다음 이동할 칸의 인덱스를 구하고, 그 칸의 내구도가 1 이상이며 로봇이 없으면 그 칸으로 이동한다. 이때 이동한 칸의 내구도를 1만큼 감소시켜야 하고 이동한 칸이 내리는 칸이면 바로 내려야 하므로 그 값은 버린다.

세 번째로 로봇 올리기 과정이다. 0번째 칸에 내구도가 1 이상이고 로봇이 없으면 새로운 로봇을 올린다. 그 칸의 내구도를 1만큼 감소시키고 robots 리스트에 0을 추가한다. 

마지막으로 내구도가 0인 칸이 K개 이상인지 확인한다. K개 이상이면 반복을 중단한다. 그렇지 않으면 단계 수를 하나씩 늘린다.

 

# PyPy3로 통과 
import sys
from collections import deque

N, K = map(int, sys.stdin.readline().split())
belt = deque(list(map(int, sys.stdin.readline().split()))) # 각 칸의 내구도
robots = [] # 로봇이 있는 칸의 인덱스

# 컨베이어 벨트의 회전
def rotate(belt, robots):
	last = belt.pop()
	belt.appendleft(last)
	# 로봇도 한 칸씩 이동 (단, 내리는 부분에 도달하면 바로 내리기)
	robots = [(robot+1)%(2*N) for robot in robots if (robot+1)%(2*N) != N-1]
	return belt, robots

# 로봇의 이동
def robot_move(belt, robots):
	for i, robot in enumerate(robots):
		next_idx = (robot+1)%(2*N)
		# 다음 칸에 로봇이 없고 내구도가 0보다 크면
		if next_idx not in robots and belt[next_idx] > 0:
			robots[i] = next_idx # 다음 칸으로 이동
			belt[next_idx] -= 1 # 다음 칸의 내구도 1 감소
    # 내리는 칸이면 내리기
	robots = [r for r in robots if r != N-1]
	return belt, robots

# 로봇 올리기
def robot_load(belt, robots, level):
	# 올리는 칸에 로봇이 없고 내구도가 0보다 크면
	if belt[0] > 0 and 0 not in robots:
		belt[0] -= 1 # 올리는 칸의 내구도 1 감소
		robots.append(0) # 로봇 올리기
	return belt, robots

# 내구도 0인 칸이 K개 이상이면 True
def end_condition(belt, K):
	if len([i for i, b in enumerate(belt) if b == 0]) >= K:
		return True
	return False


level = 1
while True:
	belt, robots = rotate(belt, robots)
	belt, robots = robot_move(belt, robots)
	belt, robots = robot_load(belt, robots, level)
	if end_condition(belt, K):
		break
	level += 1

print(level)