AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC163 E - Active Infants

問題原文

atcoder.jp

問題要旨

 N 人の幼児が横一列に並んでいる。左から  i 番目の幼児の活発度は   A_i である。
この幼児たちを1回だけ自由に並べ替える。
並び替える前左から  x 番目にいた幼児が左から  y 番目に移動した時、  abs(x - y) * A_x のうれしさが生まれる。
幼児の嬉しさの合計値を最大化せよ。

  •    1 \leq N \leq 2000

解法

幼児たちを手元に全部置いて、左 or 右に詰めていく。。。ようなことを考える。
まず、「活発度が大きい方から優先的に動かしてあげる」のは明らかに損しなさそう。ということで各幼児の最初の位置を覚えておきながら、降順ソートしておく。

あとは貪欲的にやればいいのかな〜ってなるけどどうやらサンプルみるにだめそう。   考えられるのは「各幼児を左に詰めるか右に詰めるか」で、これはナップサック問題の「使う or 使わない」と似た雰囲気を感じるのでどうやらDPで解けそうという感覚になる。

あとはどうdpテーブルを持つかだけど、「何番目までみたか」「左に何回詰めたか」「右に何回詰めたか」という情報があれば良さそう。
また、「右に何回詰めたか」というのは「何番目までみたか」「左に何回詰めたか」からわかるので、結局dpテーブルは、

 dp[i\rbrack[j\rbrack :=  i 番目の幼児まで見て、左に  j 回詰めた場合のうれしさの合計の最大値

としてもてばいい。

あとは遷移。
 dp[i\rbrack[j\rbrack からは、左に詰める or 右に詰める の2通りの遷移がある。

  • 左に詰める場合
    遷移先は  dp[i + 1\rbrack[j + 1\rbrack。比較する値は、 その幼児の最初の位置を  p として、  dp[i\rbrack[j\rbrack + (p - j) × A_i となる。

  • 右に詰める場合
    遷移先は  dp[i + 1\rbrack[j\rbrack
    比較する値は、 その幼児の最初の位置を  p とすると、すでに右に  N - (i - j) - 1 回詰めたことになるので、  dp[i\rbrack[j\rbrack + ((N - (i  - j) - 1) - p ) × A_i となる。

実装

拾うdpのほうが書きやすいかな〜と思ったけど普通に配る方が書きやすかった。
 i 人目まで見た時に左に詰めうる最大人数は  j 人なのでそこまでしかループしてない。

N = int(input())
A = list(map(int, input().split()))
A = sorted([(a, p) for p, a in enumerate(A)], reverse=True)


# dp[i][j] := i番目の幼児まで見て、左に詰める選択をj回行った場合のうれしさの最大値
dp = [[0] * (N + 1) for _ in range(N + 1)]
for i, (a, p) in enumerate(A):
    for j in range(i + 1):
        # 左に詰める場合
        dp[i + 1][j + 1] = max(dp[i + 1][j + 1], dp[i][j] + (p - j) * a)

        # 右に詰める場合
        dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + ((N - (i - j) - 1) - p) * a)

print(max(dp[-1]))

感想

本番は5分でテーブルを立てておきながら未来永劫拾う時の右詰めをバグらせ続けていた(えぇ。。。)
というか普通に配る方がデバッグも楽なんだなぁとしみじみ感じました。

とはいえJOIをやる前は手も足も出なかっただろうから、自力で黄diff近いdpを解けたのは素直に嬉しい〜〜〜!