Codeforces Round #653 (Div. 3) E2. Reading Books (hard version)

Problem - E2 - Codeforces

問題概要

$n$ 冊の本があり、本 $i$ は読むための所要時間 $t_{i}$ 、アリスがその本を好きかどうか $a_{i}$ 、ボブがその本を好きかどうか $b_{i}$ で特徴付けられる。

$n$ 冊の本からちょうど $m$ 冊の本を以下の条件を満たすように選ぶとき、その合計所要時間の最小値と、それを実現する選び方の一例を求めよ。ただし条件を満たす選び方が存在しないときは -1 を出力せよ。

選んだ本の中にアリスの好きな本が $k$ 冊以上含まれ、かつボブが好きな本が $k$ 冊以上含まれる。

制約

$1 \le k \le m \le n \le 2\times 10^{5}$
$1 \le t_{i} \le 10^{4}$

解法

まず本を $(a_{i}, b_{i})$ で分類し、 $A_{0, 0}, A_{0, 1}, A_{1, 0}, A_{1, 1}$ の4つの配列を作ります。それぞれ $t_{i}$ の値でソートしておきます。

「 $A_{1, 1}$ に属する本から少なくとも $x$ 冊選び、 $A_{0, 1}, A_{1, 0}$ に属する本から少なくとも $(k-x)$ 冊ずつ選ぶ」という値 $x$ を決め、これを $0 \le x \le k$ の範囲で全通り試します。そうすると最適な本の選び方は

$A_{1, 1}$ から、 $t_{i}$ が小さい順に $x$ 冊選ぶ。
$A_{0, 1}, A_{1, 0}$ から、 $t_{i}$ が小さい順にそれぞれ $(k-x)$ 冊選ぶ。
$r = (m-x-2(k-x))$ とする。余った本を全て混ぜ、 $t_{i}$ が小さい順に $r$ 冊選ぶ。

となります。ここで本が足りなかったり $r$ が負になるような $x$ は飛ばします。

余った本を選ぶところ以外は累積和などで計算できます。 $A_{1, 1}, A_{0, 1}, A_{1, 0}$ の中でどの本が余っているかは、 $x$ を順にループしながら追加/除外されるものを差分更新することができます。あとは余った本の $t_{i}$ を小さい方から $r$ 個選んだときの和を効率的に求められれば良いです。

これは以下のようにすればBITやセグメント木で実現できます。所要時間 $t$ をインデックスとして、以下の値を管理するBITを作ります。

$num\lbrack t \rbrack =$ 所要時間が $t$ である余り本の冊数。
$sum\lbrack t \rbrack = num\lbrack t \rbrack \times t$

小さい方から $r$ 個の和を求める時には、まず $num\lbrack t \rbrack$ の区間和を用いて二分探索することで小さい方から $r$ 番目の値を求め、 $sum\lbrack t \rbrack$ の区間和を用いてそこまでの合計を求め、取りすぎていればそのぶんを調整すれば良いです。