[BOJ] 2616 - 소형기관차 (Python)

 

 

아이디어

 

동적 계획법으로 어떤 객차까지의 객차들 중에서 최대로 운송할 수 있는 손님 수를 구해 나간다.

 

 

풀이

 

전체 객차 수를 N, 소형 기관차 한 대가 끌 수 있는 최대 객차 수를 M에 저장한다.

객차 손님 수는 앞에 0을 추가하여 리스트 acc에 저장하고, acc를 누적합 리스트로 만든다.

0으로 채워진 4*(N+1) 크기의 2차원 리스트 dp를 만들고, for문으로 진입한다.

dp[i][j]는 i 번째 소형 기관차가 j 번째 객차까지의 손님들 중 운송할 수 있는 최대 손님 수를 나타낸다.

그 다음 순번의 소형 기관차는 이전 순번의 소형 기관차가 운송하는 손님과 겹치지 않도록 M만큼씩 차이를 둔다.

 

def solution():
    N = int(input())
    acc = [0]+list(map(int, input().split()))
    M = int(input())
    for i in range(N):
        acc[i+1] += acc[i]
    dp = [[0]*(N+1) for _ in range(4)]
    for i in range(1, 4):
        for j in range(M*i, N+1):
            dp[i][j] = max(dp[i][j-1], dp[i-1][j-M]+acc[j]-acc[j-M])
    print(dp[-1][-1])

solution()

 

 

처음에는 3중 for문으로 접근했는데 시간 초과가 났다.

# 시간 초과 코드
res = 0
for i in range(N-(M-1)*3):
    a = acc[i+M]-acc[i]
    for j in range(i+M, N-(M-1)*2):
        b = acc[j+M]-acc[j]
        for k in range(j+M, N-M+1):
            c = acc[k+M]-acc[k]
            res = max(res, a+b+c)
print(res)