abc 163 e 'Active Infants' (python)

25分で d まで順調に解けたが e,f がどうにもできず終了。５完の道は遠い。
解説を読んで dp で実装してみたがなかなか手ごわかった。Unrated なので難易度がはっきりしないが本来なら青diff になっているかも。
dp[ i ][ j ] を
（An の大きいものから i 番目まで見たときに、左端(a[0]) から j 個決まっているときの最大値）
としてdp の遷移を追った
例えば n = 6 で dp[3][2] なら a[0],a[1],a[5] はすでに決まっていて、４番目に大きい数字を
左(a[2])にいれるなら dp[4][3] の候補となり、右(a[4])に いれるなら dp[4][2] の候補となる
dp[ i+1 ][ j ] の最大値の候補は、dp[ i ][ j -1 ] と dp[ i ][ j ] の２つである
初期値を小さく( -10**9*2*10**3 未満)しておかないと,
本来存在しないはずの dp[0][1] とかから dp[1][1] などが最大値を作りうるので注意
python で提出するときは def main(): をつけてないと TLE した
pypy だと 140ms くらいで通るのになぜ？

import sys
input = sys.stdin.readline
n = int(input())
a = [int(i) for i in input().split()]
b = [int(i) for i in range(n)]
ab = [(x,y) for x,y in zip(a,b)]
from operator import itemgetter
ab = tuple(sorted(ab,key=itemgetter(0),reverse=True))
dp = [[-10**14]*(n+1) for i in range(n+1)]
dp[0][0] = 0
for i in range(n):
  for j in range(i+2):
    dp[i+1][j] = max(dp[i][j]+ab[i][0]*abs(ab[i][1]-(n-1-(i-j))),dp[i][j-1]+ab[i][0]*abs(ab[i][1]+1-j))
print(max(dp[n]))