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

問題へのリンク

問題概要

正の整数からなる長さ $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;
}

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

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

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

問題概要

制約

二分探索法へ帰着

判定問題

コード