ABC162 F - Select Half - あっとのTECH LOG

問題原文

問題要旨

長さ $N$ の整数列 $A$ の中から、隣り合う部分を選ばないように $[\frac{N}{2}\rbrack$ 個選ぶ。
選んだ要素の総和を最大化せよ。

$1 \leq N \leq 10^{5}$
$1 \ leq |A| \leq 10^{9}$

解法

$dp[i\rbrack[j\rbrack :=$ $i$ 番目までみて $j$ 個選んでいる場合の最大値がぱっと思い浮かぶが、これは $O(N^{2})$ になので間に合わない。
少し考えると、「８番目まで２つ選んでいる場合〜」みたいなケースは計算しなくてよいことがわかる。
もうちょっと考えると、 $i$ 番目までみた時に、だいたい $[\frac{i}{2}\rbrack$ 個選んでおけば、最終的に $[\frac{N}{2}\rbrack$ 個選べることがわかる。
もう少し考えると、だいたい $[\frac{N}{2}\rbrack$ の ±1ぐらいの範囲だけを計算しておけばいいことがわかる。（ここらへんノリなのでちょっと自信ない）これで計算量が定数倍に３が乗る程度の $O(N)$ になったので、dpで解ける。

dpテーブルに持たせる情報は、「どこまでみたか」「何個えらんだか」「最後に選んだか」の３つ。
遷移は、

dp( $i$ 番目までみて, $x$ 個選んで, 最後に選んでない)= max(dp( $i - 1$ 番目までみて, $x$ 個選んで, 最後に選んだ), dp( $i - 1$ 番目までみて, $x$ 個選んで, 最後に選んでない))
dp( $i$ 番目までみて, $x$ 個選んで, 最後に選んだ) = max(dp( $i - 1$ 番目までみて, $x - 1$ 個選んで, 最後に選んんでない) + $A_i$

になる。

実装

こういうのは defaultdict を使うと楽。
defalt値は0でないので注意。

正直もうちょっと狭い範囲だけ計算すればいい感じがするけど、遷移元の計算が終わっていて、遷移するべきところから遷移してくれば大丈夫（たぶん）。

from collections import defaultdict
N = int(input())
A = list(map(int, input().split()))

# dp[(i, x, flg)] := i番目までみてx個選んでいる時の最大値。flgは最後の１つをとったかどうか。
dp = defaultdict(lambda: -float('inf'))
dp[(0, 0, 0)] = 0

for i, a in enumerate(A, start=1):
    j = i // 2
    for x in range(j - 1, j + 2):
        dp[(i, x, 0)] = max(dp[(i - 1, x, 0)], dp[(i - 1, x, 1)])
        dp[(i, x, 1)] = dp[(i - 1, x - 1, 0)] + a

print(max(dp[(N, N // 2, 0)], dp[(N, N // 2, 1)]))

感想

DPテーブルをdefaultdictで持つの個人的には結構好きです（気が楽なので）