[AtCoder] ABC155 D – Pairs (400点)
[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;
}