[BOJ] Q2616 소형 기관차

[BOJ] Q2616 소형 기관차

Question

Language: Python

Difficulty: Silver 2

기차 칸의 수가 최대 50000개이므로 이에 대해 3개의 그룹으로 나누는 작업으로 통해 모든 조합을 탐색하는 Brute-force 방식으로 진행하게 되면 주어진 시간에 문제를 풀이할 수 없다.

해당 문제를 실습을 예제를 통해 알아보자

기차 번호	1	2	3	4	5	6	7
0	35	40	50	10	30	45	60
1	0	35+40=75	40+50=90	50+10=60	10+30=40	30+45=75	45+60=105
2	0	0	0	75+60=135	75+60=135	60+75=135	40+105=145
3	0	0	0	0	0	135+75=210	135+105=240

각각의 기차에 대해 조사하면서 이전까지 태울 수 있는 최대 승객수를 유지하면서, 그 보다 많은 수의 승객을 태우는 경우가 존재하면 다시 최신화 시킨다.

즉, 아래의 알고리즘을 보면

Algorithm

for i in range(1,4):
  for j in range(i*max_carry,n+1):
    """
    여기서 dp[i][j-1]은 이전까지 태울 수 있는 최대의 승객수
    dp[i-1][j-max_carry] 는 해당 기차를 조사하기 이전의 기차 까지의 최대값 
    cum[j]-cum[j-max_carry]는 j-max_carry~j까지의 기차에 있는 승객수의 합을 구하기 위해 누적합 활용
    """
    dp[i][j]=max(dp[i][j-1],dp[i-1][j-max_carry]+(cum[j]-cum[j-max_carry]))

Solution

def solution():
    dp=[[0]*(num+1) for _ in range(4)]

    for i in range(1,4):
        for j in range(i*max_carry,num+1):
            dp[i][j]=max(dp[i][j-1],dp[i-1][j-max_carry]+cum_passengers[j]-cum_passengers[j-max_carry])

    return max(dp[3])   

    
if __name__ == "__main__":
    num=int(input())
    passengers=list(map(int,input().split()))
    max_carry=int(input())
    cum_passengers=[0]*(num+1)
    for i in range(1,num+1):
        cum_passengers[i]=cum_passengers[i-1]+passengers[i-1]
      
    answer=solution()
    print(answer)