백준 2616번 소형기관차: Python 풀이 및 정답 코드

문제 링크

https://www.acmicpc.net/problem/2616

2616번: 소형기관차

 

 

 

문제 설명

문제

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

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

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

예를 들어 기관차가 끌고 가던 객차가 7칸이고, 소형 기관차 1대가 최대로 끌 수 있는 객차 수는 2칸이라고 하자. 그리고 1번 부터 7번까지 각 객차에 타고 있는 손님의 수가 아래 표와 같다고 하자. 괄호속에 있는 숫자는 객차 번호를 나타낸다.

35

40

50

10

30

45

60

소형 기관차 3대는 각각 1-2번, 3-4번, 그리고 6-7번 객차를 끌고 가면 손님 240명을 운송할 수 있고, 이보다 많은 수의 손님을 운송할 수 없다.

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

입력

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

출력

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

 

 

 

 

문제 풀이

처음에 다섯시간 걸려서 이렇게 풀고 시간 초과 나왔다.

N = int(input())
numbers = list(map(int, input().split()))
M = int(input())

acc = [0]
for n in numbers:
    acc.append(acc[-1] + n)
sum_people = []
for i in range(1, N - M + 2):
    sum_people.append(acc[i + M - 1] - acc[i - 1])
for i in range(N - M + 2, N + 1):
    sum_people.append(acc[-1] - acc[i - 1])

sorted_indecis = sorted(range(len(sum_people)), key=lambda i: -sum_people[i])

dp = [[[sorted_indecis[0]]], [], []]

for i in range(1, N):
    index = sorted_indecis[i]
    dp[0].append([index])
    for nums in dp[0]:
        if not list(filter(lambda x: abs(x - index) < M, nums)):
            dp[1].append(nums + [index])
    for nums in dp[1]:
        if not list(filter(lambda x: abs(x - index) < M, nums)):
            dp[2].append(nums + [index])
    if dp[2]:
        break

print(
    max(sum(list(map(lambda x: sum_people[sorted_indecis[x]], nums))) for nums in dp[2])
)

dp에 맞는 사고과정을 하려다보니 꼬이고 꼬여서 너무 어렵게 풀어버렸다.
위 코드는 누적합을 저장한 배열을 내림차순으로 정렬하고 앞에서부터 조건에 맞게 이웃하지 않도록 세개를 고르는 코드이다.

사실 그냥 넣거나 넣지 않거나 하는 배낭 문제와 같은 방법으로 풀면 된다.

코드는 다음과 같다.

 

 

 

코드

N = int(input())
numbers = list(map(int, input().split()))
M = int(input())

acc = [0]
for n in numbers:
    acc.append(acc[-1] + n)
sum_people = []
for i in range(1, N - M + 2):
    sum_people.append(acc[i + M - 1] - acc[i - 1])
for i in range(N - M + 2, N + 1):
    sum_people.append(acc[-1] - acc[i - 1])

dp = [[sum_people[0], 0, 0]] + [[0, 0, 0] for _ in range(N - 1)]

for i in range(1, N):
    dp[i][0] = max(dp[i - 1][0], sum_people[i])
    if i - M >= 0:
        dp[i][1] = max(dp[i - M][0] + sum_people[i], dp[i - 1][1])
        if dp[i - M][1] == 0:
            dp[i][2] = dp[i - 1][2]
        else:
            dp[i][2] = max(dp[i - M][1] + sum_people[i], dp[i - 1][2])

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

 

 

 

결과

 

 

 

다른 사람의 코드

 

 

 

리뷰