ARMERIA

Rubyと競技プログラミングの話 AtCoderやCodeforcesの問題解説記事が多め。

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区間和を用いてそこまでの合計を求め、取りすぎていればそのぶんを調整すれば良いです。

これで最適値を求めることができます。この問題では具体的にどの本を選ぶかも要求されているので、最適値を実現する  x の値を記憶しておいて、改めて最適な本の選び方をなぞって選ぶ本の番号を求めれば良いです。

ACコード

Submission #85337458 - Codeforces