けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 037 C - 億マス計算 (試験管青色)

二分探索法の教育的典型問題ですね!!

問題概要

正の整数からなる長さ  N の数列が 2 つ ( a_{0}, a_{1}, \dots, a_{N-1} b_{0}, b_{1}, \dots, b_{N-1}) 与えられます。

 (i, j) に対して  a_{i} \times b_{j} を計算することで得られる  N^{2} 個の整数を考えます。

これらの整数を小さい順に並べたとき、 K 番目 (1-indexed) に来る値を求めなさい。

制約

  •  1 \le N \le 30000
  •  1 \le K \le N^{2}
  •  1 \le a_{i}, b_{j} \le 10^{9}

二分探索法へ帰着

単純に  N^{2} 個の値を計算して、それらを小さい順にソートする方法が最初に考えられるかもしれません。しかしこれでは  O(N^{2} \log N) の計算量を必要としてしまいます。

ここで「 K 番目に小さい値を求めなさい」という問題に対しては二分探索法が有効になることがある、ということを思い出してみましょう。求めたい値は、次の判定問題の答えが "Yes" となる最小の値  x に他ならないのです。


 N^{2} 個の整数のうち、 x 以下の値が  K 個以上あるかどうかを判定せよ


それではこの判定問題を解く方法を考えてみましょう。

判定問題

 (i, j) という 2 つの添字について考える問題では、どちらか 1 つの添字を固定して考えることが有効です。ここでは、 a_{i} を固定して考えてみましょう。

つまり、各  i = 0, 1, \dots, N-1 に対して、 a_{i} \times b_{j} の値が  x 以下となる  b_{j} の個数を数えていくことにします (それらの総和をとればよいです)。これは次の問題と等価です。


 b_{0}, b_{1}, \dots, b_{N-1} のうち、 \frac{x}{a_{i}} 以下であるものは何個あるか


この問題を解くためには、ふたたび二分探索 (lower_bound) が活用できます。具体的には、数列  b をあらかじめソートしておくことを前提として

upper_bound(b.begin(), b.end(), x / a[i]) - b.begin()

によって求められます。

コード

以上の解法は、次のように実装できます。

最後に計算量を見積もります。まず判定問題を解くのに要する計算量は  O(N \log N) となります。そして判定回数は、 C = \max_{i, j} a_{i} b_{j} として、 O(\log C) と評価できます。よって全体の計算量は  O(N \log N \log C) と評価できます。

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

int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> a(N), b(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    for (int i = 0; i < N; ++i) cin >> b[i];

    // b はあらかじめソートしておく
    sort(b.begin(), b.end());

    // 判定問題
    auto check = [&](long long x) -> bool {
        long long cnt = 0;
        for (int i = 0; i < N; ++i) {
            cnt += upper_bound(b.begin(), b.end(), x / a[i]) - b.begin();
        }
        return (cnt >= K);
    };

    // 二分探索 (right は十分大きな値に設定)
    long long left = 0, right = 1LL<<60;
    while (right - left > 1) {
        long long mid = (left + right) / 2;
        if (check(mid)) right = mid;
        else left = mid;
    }
    cout << right << endl;
}