ARC037 C問題億マス計算 - 輪郭をなぞるだけのブログ

昨日の ABC155 の D 問題の類題であり、D 問題はこれを引っ張ってきて何とか通せた。
実装を無限にバグらせて、通ったのコンテスト終了 10 分前くらいだけど…

atcoder.jp

問題概要

長さ $N$ の数列 $a, b$ が与えられる。
$a_{i} \times b_{j}$ の計算で現れた $N^2$ 個の値を昇順で並べたときの $K$ 番目にある値は何か。

制約

$1 \leq N \leq 30000$
$1 \leq K \leq N^{2}$
$1 \leq a_{i} \leq 10^{9}$
$1 \leq b_{j} \leq 10^{9}$

(考察 →) 実装

基本問題っぽいけど、よく分からないとなってた気がする。

まず、愚直に計算してソートは、 $O(N^2 \log N^{2})$ となって間に合わない。

そこで、答えとなる $X$ を決め打ちする。
二分探索でよくある (?) 答えの決め打ちである。
ここで、 $X$ が $K$ 番目の値になり得る条件というのは、 $X$ 以下の値が $K$ 個以上あることである、と言い換えることができる。
これを二分探索時の判定条件にする。

次に、個数をどう数えるかだが、積が $X$ 以下であるというのは、
$a_{i} \times b_{j} \leq X$ 　つまり　 $b_{j} \leq X / a_{i}$
と書くことができる。
この個数も二分探索で求めることができる。
二分探索のために、数列 $b$ をあらかじめ昇順でソートしておく。
$a_{i}$ を見ているとき、 $a_{i} \times b_{j}$ が $X$ 以下である個数は、 $X / a_{i}$ より大きくなる最初の位置 (イテレータ) を調べれば、求めることができる。
(そのイテレータを it とすると、it - b.begin() で個数は求まる。)

手順をまとめると、

数列 $b$ を昇順にソートする。
$X$ を (ng, ok] で二分探索する。このとき、数列 $a$ の要素を順に見ていき、数列 $b$ の要素との積が $X$ 以下となる個数を数え上げる。この個数を $K$ と比較して、範囲を狭めていく。

これで、計算量は、ソート部分が $O(N \log N)$ 、二分探索部分が $O(\ N \log N \times \log(max(a_{i} \times b_{j}))\ )$ となり解くことができる。

コード

横着して上限部分を $10^{18}$ にしている。

#include <iostream>
#include <vector>
#include <algorithm>
#include <cstdint>
using namespace std;

int64_t N, K;
vector<int64_t> a, b;

// X 以下の数が K 個以上か否か
// K 個以上あれば、(ng, ok] → (ng, mid]
// K 個以上なければ、(ng, ok] → (mid, ok]
bool check(int64_t X){
    int64_t cnt = 0;
    for(int i = 0; i < N; ++i){
        // a_i * b_j が X 以下になる個数を足していく
        cnt += upper_bound(b.begin(), b.end(), X / a[i]) - b.begin();
    }
    return cnt >= K;
}

int main(){
    cin >> N >> K;
    a.resize(N);
    b.resize(N);
    for(int i = 0; i < N; ++i)  cin >> a[i];
    for(int i = 0; i < N; ++i)  cin >> b[i];
    sort(b.begin(), b.end());
    
    // (ng, ok] で二分探索
    int64_t ng = 0, ok = 1e+18;
    while(ok - ng > 1){
        int64_t mid = (ng + ok) / 2;
        if(check(mid))  ok = mid;
        else    ng = mid;
    }
    cout << ok << endl;
}