[AtCoder] ABC155 D – Pairs (400点)

2020年2月17日AtCoder400点, 二分探索, K番目, 2つの積

[AtCoder] ARC037 億マス計算 の上位問題です。

問題概要

問題へのリンク

N個の整数 \(A_1, A_2, A_3, \cdots, A_n\) がある。
2つを選んでその積を \(\frac{N(N-1)}{2}\) 個小さい順に並べた時に、K 番目にある数は何か?

制約

  • \(2 \leq N \leq 2 \times 10^5\)
  • \(1 \leq K \leq \frac{N(N-1)}{2}\)
  • \(-10^9 \leq A_i \leq 10^9\ (1 \leq i \leq N)\)

考え方

単純に考えると、\(\frac{N(N-1)}{2}\) 個の積を計算してソートしなければならないので、O(N^2)以上になって間に合いません。

二分探索が基本方針

実は、K番目の数 X を求める時の典型テクニックとして、

 「小さい方からK番目の数が X」
⇔「X-1 以下の数は K個未満で、X以下の数は K 個以上」
⇔「X以下の数が K 個以上 となる最小のX」

という言い換えによって、二分探索で X を高速に求めることができます。

つまり、X を決めた時に、

  • X以下の数が K 個以上かどうか

を高速に判定することができれば二分探索が使えます。

X以下の数が K 個以上かどうかの判定

全てのペアの積を計算する時間はありません。2つのものを同時に考えないようにすると、ペアの片方を固定することを思いつきます。

ここで知りたいのは、「X以下の数が何個あるか?」です。ペアの片方を固定して \(A_i=a\) とします。単純化するために、全ての数が正の時のみを考えると、

「X以下の数が何個あるか?」
\(\Leftrightarrow \frac{1}{2} \sum_{i=1}^{N}\) (\( a\cdot b \leq X\) となる b の数(a自身は除く))
\(\Leftrightarrow \frac{1}{2} \sum_{i=1}^{N}\) (\( b \leq \lfloor \frac{X}{a} \rfloor\) となる b の数(a自身は除く))

となります。(同じペアを2回数えているので2で割っています)。

これにより、二分探索を N 回行うことで「X以下の数が何個あるか?」を求めることができます。

正負によって場合分け

先程の判定は、単純化するために、全ての数が正の時のみを考えました。しかし、制約によると入力の値は、負や0の場合も含まれています。

よって以下の3つの場合を考える必要があります。

  • X が負の時(正と負の積)
  • X が0の時(0と何かの積)
  • X が正の時 (正と正の積 or 負と負の積)

ペアの積が、負となるものの個数・0となるものの個数・正となるものの個数はそれぞれ簡単に求まります。X がどれに含まれるかは、Kの値で決まるのです。

解き方

  • minus := ペアの積が 負 となる個数
  • zero := ペアの積が 0 となる個数
  • plus := ペアの積が 正 となる個数

とすると、K番目に小さい数 X について

  • X が負の時(Kが minus 以下の時)
  • X が0の時(Kが minus+zero 以下の時)
  • X が正の時 (Kが minus+zero+plus 以下の時)

のどれになるのか判定し、X について二分探索をする。

X が負の時

正と負の積を小さい順に並べて、K番目に小さい数をXとします。

正の数 a を固定した時、負の数 b が何通りあるかを考えます。

「X以下の数が何個あるか?」
\(\Leftrightarrow \sum\) (\( a\cdot b \leq X\) となる b の数)
\(\Leftrightarrow \sum\) (\( b \leq \lfloor \frac{X}{a} \rfloor\) となる b の数)

X が 0 の時

これは0を出力して終わりです。

X が正の時

正と正の積と、負と負の積を小さい順に並べて、(K-minus-zero)番目に小さい数をXとします。

正と正の積

「X以下の数が何個あるか?」
\(\Leftrightarrow \frac{1}{2} \sum_{i=1}^{N}\) (\( a\cdot b \leq X\) となる b の数(a自身は除く))
\(\Leftrightarrow \frac{1}{2} \sum_{i=1}^{N}\) (\( b \leq \lfloor \frac{X}{a} \rfloor\) となる b の数(a自身は除く))

負と負の積

全て正に直してから、正と正の積と同様に計算すれば良いです。

C++での実装例

ABCのD問題にしては実装が重いです。