ABC438 - D - Tail of Snake

문제

뱀을 N개의 블록으로 나누고, 머리다움, 몸통다움, 꼬리다움을 평가한다. 이 평가 값을 합산할때 가장 큰 값이 나오도록 구역을 나누기로 했다.

길이가 N인 정수수열 A = ( $A_1$ , … , $A_N$ ) , B , C 가 주어질때

\[1 \le x < y < N\]

을 만족하는 정수 쌍 (x,y) 에 대해서

\[\sum_{i=1}^{x} A_i + \sum_{i=x+1}^{y} B_i + \sum_{i=y+1}^{N} C_i\]

의 최대값은?

제약조건

• $ 3 \le N \le 3 \times 10^5 $
• $ 1 \le A_i , B_i , C_i \le 10^6 $
• 입력되는 모든 값은 정수

bruteforce - 시간초과

N = 300 000 이 최대이다

N = int(input())
A = list(map(int,input().split()))
B = list(map(int,input().split()))
C = list(map(int,input().split()))

max_v = 0
for x in range(1,N):
    for y in range(x+1,N):
        s = 0
        for i in range(N):
            if i<=x:
                s += A[i]
            elif x < i and i <= y:
                s += B[i]
            else:
                s += C[i]
        max_v = max(s,max_v)
print(max_v)

이거는 N^3, 시간초과가 당연함.

dp

한번만 훑으면서 최적의 값을 저장할 수 있다

dp 설계

i번째 블록에 있을 때, 3가지 상태에 있다.

1. 상태 0 (머리 중) : 아직 머리 부분을 지나고 있다 (수열 A 선택)
2. 상태 1 (몸통 중) : 머리가 끝나고 몸통 부분을 지나고 있다 (수열 B 선택)
3. 상태 2 (꼬리 중) : 몸통이 끝나고 꼬리부분 지나는 중 (수열 C 선택)

dp[i][state] 라는 배열을 정의

• dp[i][0] : i번째 블록까지 왔을 떄, 현재 머리 상태인경우 최대 합
• dp[i][1] : 현재 몸통 상태인경우 최대합
• dp[i][2] : 현재 꼬리 상태인 경우 최대합

점화식

1. dp[i][0] (현재 머리):
   • i-1번 블록도 머리였어야 함
   • dp[i][0] = dp[i-1][0] + A[i]
2. dp[i][1] (현재 몸통):
   • 이전 블록(i-1) 에서 상태는 2가지 case
     1. 원래 몸통 -> 계속 몸통 (dp[i-1][1])
     2. 머리 -> 몸통 (dp[i-1][0])
   • max(dp[i-1][0], dp[i-1][1]) + B[i]
3. dp[i][2] (현재 꼬리):
   • 이전 블록에서 상태는 2가지 case가 있다
     1. 꼬리 -> 꼬리 (dp[i-1][2])
     2. 몸통 -> 꼬리 (dp[i-1][1])
   • max(dp[i-1][1], dp[i-1][2]) + C[i]

초기값 설정

불가능한 상태 처리 필요

\[1 \le x < y < N\]

• 최소한 1개 블록은 머리여야 하고
• 최소한 1개 몸통
• 최소한 1개 꼬리

따라서

• i=0 은 무조건 머리
• i=1 부터 몸통 등장 가능
• i=2 부터 꼬리 등장 가능

max 함수 비교시 자동 탈락되도록 -inf로 초기화 한다.

N = int(input())
A = list(map(int,input().split()))
B = list(map(int,input().split()))
C = list(map(int,input().split()))

# init
dp = [[-float('inf') for j in range(3)] for i in range(N)]
dp[0][0] = A[0]
## 나머지를 -inf로 놓으면 알아서 없어지게 됨.
## 가장 오른쪽 위의 값도 자연스럽게 안쓰이게 된다.

# dp
for i in range(1,N):
    dp[i][0] = dp[i-1][0] + A[i]
    dp[i][1] = max(dp[i-1][0], dp[i-1][1]) + B[i]
    dp[i][2] = max(dp[i-1][1], dp[i-1][2]) + C[i]

print(dp[-1][-1])

dp[i][state] 이거 떠올릴수 있으면 어떻게든 풀었을것 같은데, 여기까지 왔을때 시간이 없기도 했고 …

아직 이런 문제를 잘 못하는것 보면 좀 답답하기도 하고… 어쨌든 1년만에 atcoder를 했다