[python] 백준 2579 - 계단오르기

백준 2579 - 계단오르기 [SILVER 3]

0부터 N까지의 계단을 오르는데, 다음 칸 또는 다다음 칸으로 이동할 수 있다.

다만, 연속된 계단을 세 번 이상 밟으면 안 되며, 세 칸 이상 한번에 올라갈 수 없다.

계단마다 점수가 있으며, 이 점수를 최대화 할 수 있는 방법을 찾는다.

첫 번째 시도, BFS와 브루트 포스

n번째 계단을 밟았을 때, n-1번째 계단도 밟았다면 무조건 다음 번엔 한 칸을 건너 뛰어야 한다…

n번째 계단에 도달했을 때 최대의 값을 기록하는 것은 위의 문제로 힘들지 않을까?

경우의 수마다 각자의 값을 가지고 가는 것도 쉽지 않을 것 같다…

가지치기가 힘들 수 있으므로, 일단 브루트포스로 시작해보았다.

import sys

input = sys.stdin.readline

N = int(input())

stairs = [int(input()) for _ in range(N)]

answer_list = []

score, now_step, flag = 0,-1,0
Q = [(score, now_step, flag)] # 각 경우의 수는 누적 점수, 현재 위치, 연속으로 올라왔는지 여부를 들고 있는다.

while Q:

score, now_step, flag = Q.pop(0)
if now_step == N-1:
    answer_list.append(score)

for ns in [1, 2]:
    if flag == ns+1:
        continue
    next_step = now_step + ns
    if next_step >= N:
        break
  
    new_score = score + stairs[next_step]
    next_flag = 1 if ns == 2 else flag + 1

    Q.append((new_score, next_step, next_flag))

print(max(answer_list))

BFS로는 힘든 문제

당연히 모든 경우의 수(낮은 점수로 중복해서 방문한 것까지도)를 종합하며,

한 경우의 수 연산에 여러 개의 탈출 조건과 플래그 연산이 필요하므로 무조건 시간 초과가 발생한다.

생각해보니 이건 과거에 나를 심하게 괴롭혔던 문제다.

몇 날을 고민했는데, 초등부 정보 올림피아드 문제라는 것을 보고 깜짝 놀랐다…

두 번째 시도, DP

아무리 가지를 잘 쳐낸 BFS라고 해도 시간 초과가 날 수밖에 없는 문제고, DP를 이용한 풀이가 필요했다.

백준의 예시 문제에서 계단 점수의 배열을 arr로 지정하고 따라 올라가며 직접 계산하며,

가장 높은 경우의 값을 max라는 list에 저장한다고 생각해보자.

arr = [10, 20, 15, 25, 10, 20]
max = []

max값을 찾을 때는 현재까지 오는 데에만 조건을 만족하기만 하면 된다.
다음 칸에 대한 조건은 전혀 생각할 필요가 없다!

시작값 세팅

1은 밟거나 안 밟거나다. 당연히 밟아야 한다. max = [10]

2는 1을 밟고 올라오느냐, 한번에 올라 왔느냐다. 당연히 1+2가 높다. max = [10, 30]

3은 1을 밟고 두 칸을 뛰었는지, 2를 밟고 한 칸 올라왔는지다.

예시를 기준으로는 2 + 3이 높으므로, max = [10, 30, 35]

4번 계단

4는 어떨까?

1과 2를 밟고 두 칸을 뛰었거나, 1과 3을 밟고 한 칸 올라왔는지다.

여기서부터는 이전 경우의 수 마다의 값을 생각해야한다.

1과 2를 밟고 두 칸 오르기

1과 2를 밟았다는 것은, arr의 1번째와 2번째를 더한 값인 max의 두 번째 값, 30을 가져온다는 이야기이다.

그리고 현재 위치인 arr의 네 번째 값, 25를 거기에 더해준다.

이를 식으로 표현하면 arr[1] + arr[2] + arr[4]이며, 이는 max[2] + arr[4] 이다. (편의상 인덱스를 1부터 잡았다.)

1과 3을 밟고 한 칸 오르기

1은 arr과 max 두 리스트에서 모두 같은 값이지만, 반드시 max에서 가져와야 한다.

이전 단계의 총합을 가지고 있어야 하기 때문이다. (따라서 1번째 단계에서는 의미가 보이지 않는다.)

3은 arr에서 가져와 더해야 한다.

max에서 가져오면, arr[1] + arr[1] + arr[3]이 되므로,

나중에는 이전 단계의 총합이 두 번 더해지는 것과 마찬가지가 된다.

식으로는 max[1] + arr[3] + arr[4]이다.

이 두 개의 경우의 수에서 만들어진 값 중 높은 값을 max[4]에 지정하자.

max(max[1] + arr[3] + arr[4], max[2] + arr[4])이라는 식이 완성된다.

위의 식에서는 각각 50과 55가 나오므로, max[4]는 55가 되었다.

확실히 이해하기 위해 계산을 이어나가보자.

현재까지 max 리스트는 다음과 같다.

max = [10, 30, 35, 55]

5번 계단

5번 계단까지 오는 경우의 수는 다음과 같다.

1 - 3 - 5
1 - 2 - 4 - 5
2 - 3 - 5
2 - 4 - 5

이를 이전에 만든 식으로 풀어보면 다음과 같아진다.

arr[1] + arr[3] + arr[5] == 10 + 15 + 10 == 35
arr[1] + arr[2] + arr[4] + arr[5] == max[2] + arr[4] + arr[5] == 30 + 25 + 10 == 65
arr[2] + arr[3] + arr[5] == max[3] + arr[5] == 35 + 10 == 45
arr[2] + arr[4] + arr[5] == 20 + 25 + 10 == 55

한 칸 올라오며, 인덱스가 1 높아졌을 뿐,

보다시피 두 번째와 세 번째 경우는 이전에 세운 식과 동일한 형태임을 알 수 있다.

이를 기반으로, 현재 단계를 i로 지정한 뒤 식을 세우면 다음과 같다.

max[i-3] + arr[i-1] + arr[i] or max[i-2] + arr[i]

둘 중 하나는 경우의 수 중 최댓값을 가지고 있다.

점화식은 65를 반환하고, max[5]에 해당 값이 들어간다.

기왕 온 거 끝을 내보자.

6번 계단

1 - 2 - 4 - 6
1 - 3 - 4 - 6
1 - 3 - 5 - 6
2 - 3 - 5 - 6
2 - 4 - 6

arr[1] + arr[2] == max[2]

arr[2] + arr[3] == max[3]

arr[1] + arr[3] + arr[4] == max[4]

였음을 기억하고, 다시 식을 만들어보자.

---

arr[1] + arr[2] + arr[4] + arr[6] == max[4] + arr[6] == 75
arr[1] + arr[3] + arr[4] + arr[6] == max[4] + arr[6] == 75
arr[1] + arr[3] + arr[5] + arr[6] == 55
arr[2] + arr[3] + arr[5] + arr[6] == max[3] + arr[5] + arr[6] == 65
arr[2] + arr[4] + arr[6] == 65

이번에도 똑같이 max[3] + arr[5] + arr[6], max[4] + arr[6] 중 하나가 최대값을 갖고 있다.

이를 max[6]에 지정해주자.

6번 계단이 마지막으로, max[-1]을 출력하면 정답인 75가 출력된다.

max[3] 이후에는 arr[5]와 arr[6]이 오고, max[4] 이후에는 arr[6]이 오듯,

한 식 안의 max의 위치와 arr의 위치는 무조건 한 칸을 건너 뛰기 때문에

‘만약 max가 이전에 연속해서 올라왔으므로 한 칸을 건너 뛰어야 하는가?’를 전혀 고려할 필요가 없는 것이다.

이를 바탕으로 코드를 짜는 것은 매우 간단하다.

정답 코드

N = int(input())

arr = [int(input()) for _ in range(N)]

# N이 2이하일 경우 아래의 코드가 동작할 수 없다.
if N <= 2:
    print(sum(arr))

else:
    # 시작 값 설정
    sum_max = [arr[0], arr[0]+arr[1], max(arr[0]+arr[2], arr[1]+arr[2])]

    # sum_max 리스트의 4번째 부터 값 채우기
    for i in range(3, N):
        sum_max.append(max(sum_max[i-3] + arr[i-1] + arr[i], sum_max[i-2] + arr[i]))
    print(sum_max[-1])

결과는 다음과 같다.

메모리	시간	코드길이
31120kb	40ms	267B