ABC 163 E - Active Infants

問題

与えられた長さ $n$ の配列 $A$ に対して、順番を適当に決めたとき、$i$ の行き先を $s_i$ として $\displaystyle\sum_{i=0}^{n-1}A_i|i-s_i|$ の最大値を求めよ。($n\leq2000$)

atcoder.jp

解説

大きいものから順番に左右いずれかに配置していけばよい。
$DP[l][r]$ を区間 $[l,r)$ における最大値とする。
区間 $[l,r)$ について考えるときは大きい方から $n-(r-l)$ 個がすでに左右にうまく配置されている状態で、$A[n-(r-l)]$ を左右どちらに置くかを考える。（$A$ はあらかじめ降順にソートしておく。）
よって、DPの更新式は $a=A[n-(r-l)]$、$i$ を $a$ のもとのインデックスとして $$ \begin{aligned} DP[l][r] = \mathrm{max} \{&DP[l+1][r]+a|i-l|,\\ &DP[l][r-1]+a|i-(r-1)|\} \end{aligned} $$ となる。

実装

Pythonで出すとTLE、PyPyでACした。
メモ化再帰を使った。再帰回数の制限を増やしておかないとmaximum recursion depth exceeded in comparisonになるので注意。

import sys
sys.setrecursionlimit(10000)
n = int(input())
*A, = map(int, input().split())
A = [(a, i) for i, a in enumerate(A)]
A.sort(reverse=True)
DP = [[None for r in range(n+1)] for l in range(n+1)]

def dp(l, r):
    if l == r:
        DP[l][r] = 0
    if DP[l][r] != None:
        return DP[l][r]
    a, i = A[n-(r-l)]
    x = dp(l+1, r)+a*abs(i-l)
    y = dp(l, r-1)+a*abs(i-(r-1))
    DP[l][r] = max(x, y)
    return DP[l][r]

print(dp(0, n))

atcoder.jp

類題

区間DPの問題はDPまとめコンテストにもあった。

atcoder.jp