Codeforces Round #544 (Div. 3)E. K Balanced Teams

問題概要

人が $N$ 人おり, $i$ 番目の人の能力は $a_i$ である.

この中から $k$ 個以下のチームを作りたい. チーム作成の制約は各チーム内の任意のペアについてそのスキルの絶対値の差が $5$ 以下であることである.( $1$ 人のチームを作ることも可能である).

最大で何人をいずれかのチームに所属させることができるか？

問題へのリンク

制約

$1 \le k \le N \le 5000$
$1 \le a_i \le 10^9$

解法

制約に dp をしてくださいと書いてあるので, dp をします. $dp[ i ] [ j ] :=$ ( $i$ 番目の人までで, $j$ 個のチームを作ったときに選べる最大人数) と定義してみましょう.

まず, $i$ 番目の人をどのチームにも所属させない場合は, 定義から明らかに $dp[ i+1 ][ j ] = dp [ i ] [ j ]$ です. では, $i$ 番目の人をから新たにチームを作る場合はどうでしょうか？

元の数列をソートしておけば, $i$ 番目の人を含むチームを新しく作る場合, $a_p \le a_i + 5$ を満たす人はすべてチームに入れるのが最適です(そうしない場合, 人数が減るか, 余計なチームをさらに作ることになります). $a_p \le a_i + 5$ を満たす $p$ がどこまでかは二分探索で求められますから, そういった $p$ に対して, $dp[ p+1 ] [ j+1 ] = dp[ i ] [ j ] + (p - i + 1)$ です.

これらを実装すれば $O(Nk \mathrm{log}N)$ で答えを得ることができます. 二分探索の部分は尺取り法で実装すれば $\mathrm{log}$ を落とすことができますが, 個人的には実装が面倒なので推奨しません(二分探索でも十分高速です).

実装

Submission #145550486 - Codeforces