AtCoder Beginner Contest 155 D - Pairs

問題概要

$N (2 \leq N \leq 2 \times 10^5)$ 個の整数 $A_1, A_2, \cdots, A_N (-10^9 \leq A_i \leq 10^9)$ が与えられる．

ここから2つの整数を選んでペアにして積をとったとき，$K$ 番目に小さい値を求めよ。

解法概要

積が正，負，0になるペアの数はすぐに計算できる。$A_1, A_2, \cdots, A_N $ に含まれる正の数，負の数，0の数をそれぞれ$n_{pos}, n_{neg}, n_0$ とすると，

積が正になるペアは，同じ符号同士のペアだから $\frac{n_{pos} (n_{pos} - 1)}{2} + \frac{n_{neg} (n_{neg} - 1)}{2}$ 個
積が負になるペアは，異なる符号同士のペアだから $n_{pos} n_{neg}$ 個
積が0になるペアは，0同士のペアか0と非0のペアだから $\frac{n_{0} (n_{0} - 1)}{2} + n_{0}(n_{pos}+n_{neg})$ 個

となる。つまり， $K$ の値が $\left( n_{pos} n_{neg}, n_{pos} n_{neg}+\frac{n_{0} (n_{0} - 1)}{2} + n_{0}(n_{pos}+n_{neg}) \right] $ の範囲に入っているなら答えは 0 である。以下，0以外の値についてのみ考える。

$K$ 番目に小さい値 $ m $ というのは，$ m $ 以下の積を持つペアが $K$組存在する値 $ m $ であると言い換えることができる (等しい積を持つペアが複数あっても適当に順序づけするとする) 。$ m $ を増やしたときに $ m $ 以下の積を持つペアが減ることはないので，「$ m $ 以下の積を持つペアが $K$ 組以上存在するか？ ($ m $ は小さい順に並べて $K$ 番目以降か？) 」の境目を二分探索で求めることができる。

あとは二分探索の判定パートを $O(n)$ や $O(n \log n)$ あたりで書ければ解ける。愚直にやると判定パートだけで $O(n^2)$ になってしまうので，あらかじめ値をソートしておいて，判定パートでも二分探索を使う (ソースコード中の，for の中で lower_bound や upper_bound を使っている部分にあたる)。

ペアの片方の値を固定して積が mid 以下になるようなもう片方の値の数を数えて足しこんでいく。私の実装では関数 is_ok の中で答えが正か負かによって場合分けをしている。この場合分けは，$K$ の値が $n_{pos} n_{neg}$ 以下かどうかで決まる。

答えが負の場合はペアのうち正の値側を固定し負の側は積が mid 以上となるものを数え (なので，切り上げ除算を使っている) ，全体の数から引く。答えが正の場合はペアが両方同じ符号なので，絶対値の小さい側を固定し，積が mid 以上となるような大きい側の値の数を数えている。2つの整数を選ぶときの順序は不問なので，同じペアを2回数えないように注意。

ソースコード

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

using ll =  long long;

constexpr ll LINF = 1LL << 60;

bool is_zero(ll k, ll n_pos, ll n_neg, ll n_zero) {
    return n_pos * n_neg < k && k <= n_pos * n_neg + n_zero * (n_pos+n_neg) + n_zero * (n_zero-1) / 2;
}

bool is_ok(ll mid, ll k, ll n_pos, ll n_neg, ll n_zero, vector<ll> &pos, vector<ll> &neg) {
    ll kk = 0;  // mid 未満の積をもつ組数
    if(k <= n_pos * n_neg) {
        for(int i=0;i<n_pos;++i) {
            ll d = n_neg - ll(lower_bound(neg.begin(), neg.end(), (-mid+pos[i]-1)/pos[i]) - neg.begin());
            kk += d;
        }
    } else {
        kk = n_pos * n_neg + n_zero * (n_pos+n_neg) + n_zero * (n_zero-1) / 2;
        for(int i=0;i<n_pos;++i) {
            ll d = ll(upper_bound(pos.begin()+i+1, pos.end(), mid/pos[i]) - pos.begin()) - (i+1);
            kk += d;
        }
        for(int i=0;i<n_neg;++i) {
            ll d = ll(upper_bound(neg.begin()+i+1, neg.end(), mid/neg[i]) - neg.begin()) - (i+1);
            kk += d;
        }
    }

    return kk >= k;
}

int main() {
    ll n, k, a;
    cin >> n >> k;
    vector<ll> pos, neg;
    for(int i=0;i<n;++i) {
        cin >> a;
        if(a > 0) pos.push_back(a);
        else if(a < 0) neg.push_back(-a);
    }
    sort(pos.begin(), pos.end());
    sort(neg.begin(), neg.end());
    ll n_pos = pos.size(), n_neg = neg.size(), n_zero = n - n_pos - n_neg;

    if(is_zero(k, n_pos, n_neg, n_zero)) {
        cout << 0 << endl;
        return 0;
    }

    // 「小さい順から k 番目以下」を取るためには，積の値は mid で足りるのか？
    ll ng = -LINF, ok = LINF;
    while(ok-ng > 1LL) {
        ll mid = (ok + ng) / 2;
        if(is_ok(mid, k, n_pos, n_neg, n_zero, pos, neg)) {
            ok = mid;
        } else {
            ng = mid;
        }
    }

    cout << ok << endl;

    return 0;
}

所感

なんてことはない (と思いたい) 少し実装が重いだけの二分探索典型問題だと思うんですが，本番で WA 連発して破滅した反省の意を込めて……。