Official

E - Amusement Park Editorial by sugarrr


「楽しさ」\(0\) 以下のアトラクションに乗ることは考えなくてよいです。

よって、問題は次のように言い換えられます。

数列 \(B=(1,2,\dots A_1,1,2,\dots A_2, \dots, 1,2,\dots A_N)\) から \(K\) 個以下の要素を選ぶとき、合計の最大値はいくつか?

数列 \(B\) の長さが \(K\) 以下の時は単純に全ての合計が答えです。
以下では数列 \(B\) の長さは \(K\) を超えるとします。

\(B\) をソートして大きい順に \(K\) 個選べば良いですが、\(A_i,K\) が大きいので工夫する必要があります。

二分探索を用いましょう。以下の判定問題を考えます。

数列 \(B\) に含まれる \(m\) 以上の値は \(K\) 個以下か? \(\cdots(*)\)

\((*)\)\(m=1\) のとき \(False\) で、 \(m=2 \times10^9+1\) のとき \(True\) です。
また、\((*)\)\(m\) に関して単調です。
さらに、\((*)\)\(O(N)\) で解くことができます。

よって、 \(m\) に対して二分探索を使うことができ、\((*)\) を満たす最小の \(m\) を求めることができました。 (求めた最小の \(m\)\(M\) とします。)

数列 \(B\) から \(K\) 個の要素を選ぶとき、

  • \(M\) 以上の要素を全て選ぶ
  • 残り(\(K\) 個に満たない部分)は \(M-1\) を選ぶ

という選び方が明らかに最適です。

以上より、この問題を解くことができました。
計算量は \(O(N \log \max(A_i))\) です。

posted:
last update: