[백준] 2616 : 소형기관차 / python

1. 문제 설명

2616번: 소형기관차

 

기차는 맨 앞에 있는 기관차 1대가 손님이 탄 객차 여러 칸을 끌고 간다.
기관차가 고장나면 기차를 운행할 수 없게 되므로
최근 철도청은 기관차 고장에 대비하여 몇몇 역에 소형 기관차 3대를 배치하기로 결정하였다.
소형 기관차는 평소에 이용하는 기관차보다 훨씬 적은 수의 객차만을 끌 수 있다.

 

기관차가 고장났을 때 끌고 가던 객차 모두를 소형 기관차 3대가 나누어 끌 수 없기 때문에,
소형 기관차들이 어떤 객차들을 끌고 가는 것이 좋을까하는 문제를 고민하다가
다음과 같이 하기로 결정하였다.

1. 소형 기관차가 최대로 끌 수 있는 객차의 수를 미리 정해 놓고,
   그보다 많은 수의 객차를 절대로 끌게 하지 않는다.
   3대의 소형 기관차가 최대로 끌 수 있는 객차의 수는 서로 같다.
2. 소형 기관차 3대를 이용하여 최대한 많은 손님을 목적지까지 운송하도록 한다.
   각 객차 마다 타고 있는 손님의 수는 미리 알고 있고, 다른 객차로 손님들이 이동하는 것은 허용하지 않는다.
3. 각 소형 기관차는 번호가 연속적으로 이어진 객차를 끌게 한다.
   객차는 기관차 바로 뒤에 있는 객차부터 시작하여 1번 부터 차례로 번호가 붙어있다.

기관차가 끌고 가던 객차의 수와 각 객차에 타고 있던 손님의 수,
그리고 소형 기관차가 최대로 끌수 있는 객차의 수가 주어질 때,
소형 기관차 3대를 이용하여 최대로 운송할 수 있는 손님 수를 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 기관차가 끌고 가던 객차의 수가 입력된다. 그 수는 50,000 이하이다.
둘째 줄에는 기관차가 끌고 가던 객차에 타고 있는 손님의 수가 1번 객차부터 차례로 입력된다.
한 객차에 타고 있는 손님의 수는 100명 이하이고, 입력되는 숫자들 사이에 빈칸이 하나씩 있다.
셋째 줄에는 소형 기관차가 최대로 끌 수 있는 객차의 수가 입력된다.
그 수는 기관차가 끌고 가던 객차 수의 1/3보다 적다.

 

출력

한 줄에 소형 기관차 3대를 이용하여 최대로 운송할 수 있는 손님 수를 출력한다.

 

2. 풀이 전 계획과 생각

✏️ 필요한 변수

• N → 기관차가 끌고 가던 객차의 수
• people → 객차에 타고 있는 손님의 수
• standard → 최대로 끌 수 있는 객차의 수

 

1️⃣ 백트레킹 + 누적합

최대로 끌 수 있는 객차의 수가 정해져 있다.

소형기관차 3개가 끌고 갈 수 있는 객차의 경우의 수를 다 구해주고

그 중 최대로 운송할 수 있는 손님 수를 갱신한다.

최대로 운송할 수 있는 손님 수를 출력해주자.

 

1) 소형기관차 3개가 연속으로 끌고 갈 수 있는 객차의 경우의 수 구하기
idx 0~N-1에서 겹치지 않고 연속되면서 최대로 끌 수 있는 객차의 수를 초과되지 않도록 3개의 구간을 뽑아주자.

경우의 수가 최소 50000C3 이상 넘는다.
⇒ 시간 초과❗️ ⏰

 

경우의 수는 객차를 선택한 경우, 선택하지 않는 경우 2가지가 있다. 
이 큰 2가지의 경우 중 가지치기를 해보자.

def go(idx, small_train_num, selected_coach_num, is_selected, passenger_num):
	# 모든 객차를 탐색한 경우
	if idx >= total_coach_num:
		return

	# 해당 객차를 선택한 경우
	if is_selected:
		# 선택한 객차 수가 최대 객차 수를 넘으면 0으로 초기화 한다.
		if selected_coach_num >= max_coach_num:
			selected_coach_num = 0
		# 새로운 객차를 선택하는 경우 새로운 소형 기관차를 준비한다.
		if selected_coach_num:
			small_train_num += 1
		# 소형기관차가 4개 이상인 경우
		if small_train_num >= 4:
			return
		
		# 총 승객 수를 누적한다.
		pessenger_num += passengers[idx]
		# 최대 승객 수를 갱신한다.
		max_passenger = max(passenger_num, max_passenger)
		
	else:
		# 해당 객차를 선택하지 않으면 소형 기관차의 객체가 끊기므로 초기화
		selected_coach_num = 0

	go(idx + 1, small_train_num, selected_coach_num, True, passenger_num)
	go(idx + 1, small_train_num, selected_coach_num, False, passenger_num)

go(0, 0, 0, True, 0)
go(0, 0, 0, False, 0)

2) 최대로 운송할 수 있는 손님 수 구하기 뽑았다면, 누적합을 통해 합을 O(1)로 구해주자.
(연속으로 끌고 갈 수 있는 객차이기에 합을 일일이 더하면서 찾아주지 말고 누적합을 통해 구할 수 있다.)

 

3. 풀이하면서 막혔던 점과 고민

🤔 백트레킹 + 누적합으로 접근하니 시간초과가 발생 어떻게 문제를 해결해줄 수가 있을까?

모든 경우를 다 따지면 시간초과가 나고,
소형기관차는 연속된 객차를 선택해야 하므로 객차를 정렬 시킬 수도 없음으로 최적해로도 풀 수가 없다.

그렇다면, DP로 접근해보자.

 

2️⃣ DP

소형기관차는 3개이며 객차를 고르는 문제이므로, Knapsack Problem로 접근해보자. (소형기관차의 순서는 필요없다.)

그러므로, 문제를 다음과 같이 정의해줄 수 있다.

 

문제 정의

dp[i][j] = 소형기관차 i개까지 선택하고, j번째까지 객차를 고려했을 때의 최대 승객의 수로 정의

 

구하고자 하는 바

dp[3][N] = 소형기관차 3개까지 선택하고, N번째(= 마지막)까지 객차를 고려했을 때의 최대 승객의 수

 

점화식

경우가 2가지가 있다.

• 현재 소형기관차가 추가 되어 j번째 객차를 추가한 경우
• 현재 소형기관차가 추가 되어 j번째 객차를 추가하지 않는 경우
• dp[i][j] = dp[i-1][j-m] + sum[j] - sum[j - m] (단, j ≥ m) (소형 기관차는 최대 m개의 객실을 끌 수 있으므로 j - m + 1 번 부터 j번 객실까지 운송)
• dp[i][j] = dp[i][j-1]

⇒ dp[i][j] = max( dp[i-1][j-m] + sum[j] - sum[j - m], dp[i][j-1] )

 

Base 조건

dp[i][j] ⇒ 0으로 초기화

 

4. 풀이

N = int(input())
people = list(map(int, input().split()))
people.insert(0, 0)

standard = int(input())

dp = [[0] * (N + 1) for _ in range(4)]

prefix_sum = [0]
for i in range(1, N + 1):
    prefix_sum.append(prefix_sum[i - 1] + people[i])

for i in range(1, 4):
    for j in range(standard, N + 1):
        dp[i][j] = max(dp[i][j - 1], dp[i - 1][j - standard] + prefix_sum[j] - prefix_sum[j - standard])

print(dp[3][N])

 

5. 풀이 후 알게된 개념과 소감

• 입력값이 5 * 10 ^ 4이므로, dp와 누적합으로 시간을 줄여 주었다.
• 객차를 골라야하는 문제이므로, Knapsack Problem으로 접근을 해서 해결했다.