[백준_20055] 컨베이어 벨트 위의 로봇 python

# n, k 입력
n, k = map(int, input().split())

# num_list 입력
num_list = list(map(int, input().split()))

# up 설계 [내구도, 로봇 유무]
up = [
    [0 for _ in range(2)]
    for _ in range(n)
]
for i in range(n):
    up[i][0] = num_list[i]

# down 설계
down = []
for i in range(2*n-1, n-1, -1):
    down.append(num_list[i])

# 함수들
# zero_more_than_k()
def zero_more_than_k():
    
    # 내구도 0인 칸의 개수
    zero_cnt = 0

    for i in range(n):
        # up의 내구도가 0이면,
        if up[i][0] == 0:
            # 칸의 개수의 추가
            zero_cnt += 1
        # down의 내구도가 0이면,
        if down[i] == 0:
            # 칸의 개수에 추가
            zero_cnt += 1
    
    # 내구도 0인 칸의 개수가 k개 이상인지 반환
    return zero_cnt >= k

# robot_on_drop_pos()
def robot_on_drop_pos():
    
    # up의 마지막칸에 로봇이 있는지 반환
    return up[-1][1]

# drop_robot()
def drop_robot():
    
    # 전역 변수 선언
    global up
    
    # 로봇 내려주기
    up[-1][1] = 0

# circulate()
def circulate():
    
    # 전역 변수 선언
    global up, down
    
    # temp_up, temp_down 설정
    temp_up, temp_down = up[-1][0], down[0]

    # up 돌려주기
    for i in range(n-1, 0, -1):
        up[i] = up[i-1]
    up[0] = [temp_down, 0]

    # down 돌려주기
    for i in range(n-1):
        down[i] = down[i+1]
    down[-1] = temp_up

# can_move(idx)
def can_move(idx):
    
    # 현재 칸에 로봇이 있고, 다음 칸에 로봇이 없고, 다음 칸의 내구도가 1 이상 남아있으면 가능
    return up[idx][1] == 1 and up[idx+1][1] == 0 and up[idx+1][0] >= 1

# move(idx)
def move(idx):
    
    # 전역 변수 선언
    global up

    # 현재 로봇을 없애주고
    up[idx][1] = 0
    # 다음 위치에 로봇을 넣어주고
    up[idx+1][1] = 1
    # 다음 위치의 내구도 내려주기
    up[idx+1][0] -= 1

# put_pos_ok()
def put_pos_ok():
    
    # up의 맨 앞자리의 내구도가 1 이상인지 확인
    return up[0][0] >= 1

# put_robot()
def put_robot():
    
    # 전역 변수 선언
    global up

    # 놓는 위치에 로봇 놓아주기
    up[0][1] = 1
    # 놓는 위치의 내구도 내려주기
    up[0][0] -= 1

# 설계
# 단계
cnt = 0

while True:

    # 내구도가 0인칸이 k개 이상이면
    if zero_more_than_k():
        # 단계 출력 후
        print(cnt)
        # 종료
        break

    # 내리는 위치에 로봇이 있다면,
    if robot_on_drop_pos():
        # 로봇 내려주기
        drop_robot()
    
    # 컨베이어 벨트 회전
    circulate()

    # 내리는 위치에 로봇이 있다면,
    if robot_on_drop_pos():
        # 로봇 내려주기
        drop_robot()
    
    # up 컨베이어 벨트에서
    for i in range(n-2, -1, -1):
        # 이동할 수 있는 로봇이면,
        if can_move(i):
            # 이동시켜주기
            move(i)
    
    # 내리는 위치에 로봇이 있다면,
    if robot_on_drop_pos():
        # 로봇 내려주기
        drop_robot()
    
    # 올리는 위치에 칸의 내구도가 0이 아니면
    if put_pos_ok():
        # 로봇 올려주기
        put_robot()
    
    # cnt 올려주기
    cnt += 1