ABC267 D - Index × A(Not Continuous ver.) Python解答例

AtCoder Beginner Contest 267 D - Index × A(Not Continuous ver.)をPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

長さ$N$の整数列$A=(A_1,A_2,\dots,A_N)$が与えられます。

長さ$M$の$A$の部分列(連続でなくてもよい)$B=(B_1,B_2,\dots,B_M)$に対する、$\displaystyle \sum_{i=1}^{M} i \times B_i$の最大値を求めてください。

注記

数列の部分列とは、数列から$0$個以上の要素を取り除いた後、残りの要素を元の順序で連結して得られる数列のことを言います。

例えば、$(10, 30)$は$(10, 20, 30)$の部分列ですが、$(20, 10)$は$(10, 20, 30)$の部分列ではありません。

制約

• $1 \le M \le N \le 2000$
• $- 2 \times 10^5 \le A_i \le 2 \times 10^5$
• 入力は全て整数。

解答例

動的計画法

整数列$A$の要素から$M$個選ぶ場合の数は${}_N C_M$ありますが、これらを全探索すると当然間に合いません。

この問題で求める値(以下、スコアと呼びます)$\sum^{M}_{i = 1} i \times B_i$を見ると、$B$の要素を$j - 1$個選んだ状態から要素を更に1つ付け加えるとき、加算される値は$j \times B_i$であり、それまで選んだ値に依存しないことがわかります。

動的計画法で解けそうな気がしてきました。

DPテーブルの定義

DPテーブルを以下のように定義します。

$dp[i][j]$：整数列$A$の$i$番目の要素までを選択対象とし、その中から要素を$j$個選んで整数列$B$を作ったときの、スコアの最大値($1 \leq i \leq N, 0 \leq j \leq M$)

初期値は以下のようになります。要素を1つも選ばなかったときスコアは0です。また、今回の問題では負の値を持つ要素が$A$の中にある(スコアは負の値を取る可能性がある)ので、初期値として十分に小さい負の値を入れておきます。

$$dp[i][j] = \begin{cases} 0 &\quad (j = 0) \\ -\infty &\quad (1 \leq j \leq M) \end{cases}$$

DPテーブルの遷移は以下のようになります。

• $i$番目の要素を選ばなかったとき：$dp[i][j] = dp[i - 1][j]$
• $i$番目の要素を選ぶとき：$dp[i][j] = \max(dp[i][j], dp[i - 1][j - 1] + j \times A_i)$

遷移についての補足

$i$番目の要素を選ぶとき、それを選ぶことによって選んだ要素の数が$j$個になるので、「$i - 1$番目の要素までを選択対象として$j - 1$個選んだ」ときのスコアに$j \times A_i$が加算されることになります。

また、$i$番目の要素を選ばなかったとき、選んだ要素の個数は$j$個のままなので、$dp[i - 1][j]$をそのまま引き継ぎます。

答えは$dp[N][M]$となります。

実装例

Pythonによる実装例を以下に示します。

d.py

from typing import *


def main():
    # 入力受け取り
    N, M = map(int, input().split())
    A: List[int] = list(map(int, input().split()))

    # DPテーブル
    dp: List[List[int]] = [[-(1 << 60) for j in range(M + 1)] for i in range(N + 1)]
    # 初期条件
    for i in range(N + 1):
        dp[i][0] = 0

    # 遷移
    for i in range(1, N + 1):
        for j in range(1, M + 1):
            # i番目の要素を選ばないとき
            dp[i][j] = dp[i - 1][j]
            # i番目の要素を選ぶとき
            dp[i][j] = max(dp[i][j], dp[i - 1][j - 1] + j * A[i - 1])

    # 答えの出力
    print(dp[N][M])


if __name__ == "__main__":
    main()

実際の提出はこちら。