Codeforces Round #521 (Div. 3)F2. Pictures with Kittens (hard version)

問題概要

長さ $n$ の整数列 $a$ が与えられる. この中からちょうど $x$ 個の整数を選択するが, その時少なくとも連続する $k$ 個の中の一つが選ばれてなければならない. この時, 選んだ整数の合計の最大値はいくらか.

制約

$1 \le k, x \le n \le 5000$

解法

dp をしましょう. $dp[ i ] [ j ] :=$ ( $i$ 個の整数を, $j$ 番目までの整数で選び, かつ $a_j$ を選んだ時の合計の最大値 )と定義します. 通常の dp では添字が普通逆ですが, 逆だと後々困るのでこう定義しておきます.

遷移を考えます. $j$ 番目を選ぶ時, 一個前に選んだ整数は $k$ 個以内でなければなりませんから,

$dp[i+1][j+1] := \mathrm{max}(dp[i][j-k+1], dp[i][j-k+2], \cdots, dp[i][j]) + a_j$

のような遷移式ができます. 最後に選ぶ整数は $a_{n-k+1}, a_{n-k+2}, \cdots, a_n$ のいずれかですから, 答えは $\mathrm{max}_{n-k+1 \le j \le n}dp[x][j]$ となります.

これを愚直に実装すると計算量は $O(nkx)$ なので, これだと TLE します(F1 は通ります). なんとかして高速化しましょう.

dp の遷移式を見れば, $i$ についての添字は一つ前しか影響せず, $j$ についてはある範囲の最大値となっています. 典型的なテクニックですが, dp テーブルを SegmentTree として持つことにより区間最大値を高速に取得できるので, 一度の遷移が $O(\mathrm{log}k)$ で済みます. この時の全体の計算量は $O(nx\mathrm{log}k)$ です.

ただし, 添字 $i$ をあらわに持つと MLE します. in-place 化をすることで対処しましょう(一個先を作成して元の配列と swap することを繰り返すことでメモリを抑えるテクニックです).

さらに in-place 化してもなんか TLE します. 定数倍高速化のため, 更新する $j$ の範囲を限定すれば間に合わせることができます. 例えば, $dp[5][4]$ などは, 定義に合わせると $5$ 個の整数を $4$ 番目までで選ぶことに対応しているため, 明らかに値をとらず更新が不要です.

ACL のSegTree など高速な SegmentTree を用いないとどうやっても間に合わないかもしれません. 公式解説にはさらに $\mathrm{log}$ を落とした方法が載っているので確かめてみてください.

実装

Submission #146505254 - Codeforces