きなこの精進日記[python]

コドフォ Educational Contest 92 B Array Walk 解説[python]

概要

正の整数列aが左から順番に並んでおり、現在先頭の要素に存在する。

各stepで現在の位置から左右の要素に移動することができ、移動した場所の要素の値だけ得点が入る。

 動き方の制約として、連続して左の要素にいくことはできない、左にいくのは最大でz 回まで、という条件がある時、k回移動した時の得点の最大値を求めよ

制約

n <= 10**5
a <= 10**4
k <= n-1
z <=min(5,k)

考察

動き方が2種類しかないので最初全列挙を考えたが流石にTLE。

 現在位置と今何回左に進んだか、が同じ時、今までの動き方によらず、後の理想的な動き方は同じになるので、状態をまとめることができる。

 状態といて、現在位置 i と今何回左に進んだか j を持つと、現在いる位置(要素)は i - 2* j で計算できる。
 
 後持つべき状態としては、連続して左にいけないという制約から、前回左に進んだかどうかをもてば良い。

実装

3要素を持つdpを作るか、前回左に進んだかどうかでdpを分けるとかける。

import sys
import math

read = sys.stdin.buffer.read
# input = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
t = int(input())
for case in range(t):
    n, k, z = map(int, input().split())
    a = list(map(int, input().split()))
    # dp[i][j]= 現在 i 回移動して、その中でk 回左に移動した時の最大値
    INF = 10 ** 4
    dp_left_ok = [[-INF] * (1 + z) for _ in range(1 + k)]
    dp_left_ng = [[-INF] * (1 + z) for _ in range(1 + k)]

    dp_left_ng[0][0] = a[0]
    dp_left_ok[0][0] = -INF  # 最初はだめ
    # pos = i - 2*j
    for j in range(z + 1):
        for i in range(1, k + 1):
            if j > 0 and (i - 1) - 2 * (j - 1) >= 0:  # 現在地が0以上の時のみ使用
                dp_left_ng[i][j] = dp_left_ok[i - 1][j - 1] + a[i - 2 * j]
            if (i - 1) - 2 * j >= 0:
                dp_left_ok[i][j] = (
                    max(dp_left_ok[i - 1][j], dp_left_ng[i - 1][j]) + a[i - 2 * j]
                )
    print(max(max(dp_left_ng[k]), max(dp_left_ok[k])))
    """print(dp_left_ng)
    print(dp_left_ok)"""