Codeforces Round #629 (Div. 3) F. Make k Equal

お題箱より。

Problem - F - Codeforces

リクエストに沿って、正当性の説明をします。ちゃんと証明を組むとなかなか大変でした…

問題概要

$n$ 要素の整数列 $a$ と整数 $k$ が与えられる。

1回の操作で、この数列に以下のいずれか1つを行うことができる。

操作 $\min$ ： $a$ の最小要素を1つ選び、その値を1増やす。
操作 $\max$ ： $a$ の最大要素を1つ選び、その値を1減らす。

$a$ の中にある整数値が $k$ 個以上含まれているようにしたい。そのために必要な最小の操作回数を求めよ。

制約

$1 \le k \le n \le 2\times 10^{5}$
$1 \le a_{i} \le 10^{9}$

解法

まず、初期状態で $a$ の中に同じ値が $k$ 個以上含まれている場合、答えは $0$ です。以降はそうでない場合を仮定します。また $a$ はソートされているものとします。

いきなり最適な操作を考えるのは難しいので、いくつか決め打ちをします。

「最終的に値 $T$ を $k$ 個集める」という目標値 $T$ 。
「操作 $\min$ / 操作 $\max$ のうち、こちら側は必ず行うことにする」という操作。

解説上は操作 $\min$ を選んだと仮定します。

操作 $\min$ で値 $T$ の要素を増やそうとすると、まずは $a$ の最小値を $T-1$ にしなければなりません。そのためまず行うことは、「 $a_{i} \lt T$ である要素が全て $T-1$ になるまで操作 $\min$ を行う」ことです。

この $T-1$ である要素は、さらに1回の操作で値を $T$ にできます。これは値 $T$ の要素を増やすための最も「効率の良い」操作なので優先的に行うべきです。つまり、

$a_{i} \le T$ である要素が $k$ 個以上ある場合、 $k$ 個になるまで $T-1$ を $T$ に変え、操作を終了する。
$a_{i} \le T$ である要素が $k$ 個未満である場合、それら全てを $T$ にしてしまってから、残りを操作 $\max$ で調達する。

というのが最適になります。後者の場合の $\max$ 操作は $\min$ 操作と同様に、まず最大値が $T+1$ になるまで操作を行ってから、必要なだけ $T$ に変えていくことになります。

この2通りのパターンを図示します。

f:id:betrue12:20200416231229p:plain

この矢印の長さの合計が操作回数に対応します。こうして見ると、この図のように「最初からその値である要素が存在しないような値」を $T$ としてしまった場合、その上下どちらかにある「最初から存在する値」に目標値を変更することで、それより操作回数を少なく（もしくは変わらないように）できることが分かります。つまり目標値として考える必要があるのは、最初から存在している高々 $n$ 通りの値だけで十分です。

あとはそれぞれの値を目標値 $T$ とした時の答えを求めれば良いです。 $a$ をソートしておけば、 $a_{i} \le T$ である要素の個数は二分探索などで求められます。そうすると先ほどの図のように最終的に作るべき数列の形が分かるので、累積和などを用いてその形にするための合計操作回数を求めることができます。

最初に決め打ちした、「こちら側は必ず行うことにする」とした操作が $\max$ 操作である場合の最適値も同様に求める必要があります。実装テクニックとして、例えば全要素を $-1$ 倍して再ソートあるいは反転させることで同じ処理を使い回すことができます。