[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問題にしては実装が重いです。

#include <bits/stdc++.h>
#define LOOP(n) for (int _i = 0; _i < (n); _i++)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define FOR(i, r, n) for (int i = (r); i < (n); ++i)
#define All(obj) begin(obj), end(obj)
using namespace std;
using ll = long long;
using ull = unsigned long long;

int N;
ll K;
vector<ll> P, M, Z;

bool judge_m(ll mid) {
    ll sum = 0;
    for (int i = 0; i < (int)P.size(); i++) {
        ll d = (mid - P[i] + 1) / P[i];
        sum += upper_bound(M.begin(), M.end(), d) - M.begin();
    }
    return sum >= K;
}

bool judge_p(ll mid) {
    ll sum = 0;
    for (int i = 0; i < (int)P.size(); i++) {  // 正と正の積
        ll d = mid / P[i];
        ll tmp = upper_bound(P.begin(), P.end(), d) - P.begin();
        if (i < tmp) tmp--;  // 自身を除く
        sum += tmp;
    }
    for (int i = 0; i < (int)M.size(); i++) {  // 負と負の積
        ll d = mid / M[i];
        ll tmp = upper_bound(M.begin(), M.end(), d) - M.begin();
        if (i < tmp) tmp--;  // 自身を除く
        sum += tmp;
    }
    sum /= 2;
    return sum >= K;
}

int main() {
    cin >> N >> K;
    vector<ll> A(N);
    REP(i, N) cin >> A.at(i);
    REP(i, N) {  // Aを正か負か0かで分類
        if (A[i] > 0) {
            P.push_back(A[i]);
        } else if (A[i] < 0) {
            M.push_back(A[i]);
        } else {
            Z.push_back(A[i]);
        }
    }

    ll p = P.size();
    ll m = M.size();
    ll z = Z.size();
    ll minus = p * m;
    ll zero = z * (z - 1) / 2 + z * p + z * m;
    ll plus = p*(p-1)/2 + m*(m-1)/2;

    if (K <= minus) {  // X が負の時
        sort(All(M));
        ll l = -1e18 - 1, u = -1;  // 二分探索: (l,u]に答えがある
        while (u - l > 1) {
            ll mid = (u + l) / 2;
            if (judge_m(mid)) {
                u = mid;
            } else {
                l = mid;
            }
        }
        cout << u << endl;

    } else if (K <= minus + zero) {  // X が 0 の時
        cout << 0 << endl;

    } else if (K <= minus + zero + plus) {  // X が正の時 (else if でなくて else でも良い)
        K -= (minus + zero);
        REP(i, m) M[i] = abs(M[i]);  // 中身を正に変更
        sort(All(P));
        sort(All(M));
        ll l = -1, u = 1e18;  //二分探索: (l,u]に答えがある
        while (u - l > 1) {
            ll mid = (u + l) / 2;
            if (judge_p(mid)) {
                u = mid;
            } else {
                l = mid;
            }
        }
        cout << u << endl;
    }
    return 0;
}