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

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

競プロ典型 90 問 007 - CP Classes(★3)

とても教育的な二分探索の問題ですね!

この問題は、より高度な問題では部分的に何度も登場するような、極めて典型的な問題なので、そのまま覚えてしまうくらいでいいと思います!!!

問題概要

 N 個の整数  A_{0}, A_{1}, \dots, A_{N-1} が与えられます。

この数列に対して  Q 回のクエリが与えられます。各クエリでは整数  B が与えられて、次の問いに答えてください。


 i i = 0, 1, \dots, N-1 の範囲で動かしたときの、 |A_{i} - B| の最小値を求めてください


制約

 1 \le N, Q \le 300000

方針

各クエリに対して、 i = 0, 1, \dots, N-1 の場合を全探索するのでは、クエリごとに  O(N) の計算量を要するため、全体で  O(NQ) の計算量となります。これでは間に合いません。

そこで、二分探索を活用することで、各クエリに要する計算量を  O(\log N) にします。

まず、あらかじめ  N 個の整数  A_{0}, A_{1}, \dots, A_{N-1} を小さい順にソートしておきます。その後、各クエリに対して次の手順で求められます。


  1.  A_{j} \ge B となる最小の  j を求める (計算量: O(\log N))
  2. このとき、 B に最も近いのは  A_{j} A_{j-1} のうちのどちらかなので、比較して小さい方を答える (計算量: O(1))

このうちの 1. のパートについて、詳しく見ていきます。

 

 A_{j} \ge B となる最小の  j を求める

この処理は実は、C++ では標準ライブラリ std::lower_bound()、Python3 でも標準ライブラリ bisect_left() として用意されているくらい典型的な処理です。

処理内容は下図のような感じです。つまり、 A を小さい順に見て行ったとき、初めて  B に追いつく瞬間を捉える処理です。これらはいずれも  O(\log N) の計算量で実行できます。

C++ の場合

C++ であれば、

vector<int>::iterator it = lower_bound(A.begin(), A.end(), B);

によって、そのような  Aイテレータが取得できます。添字  A_{j} \ge B となる最小の添字  j を取得したいときは、

int j = lower_bound(A.begin(), A.end(), B) - A.begin();

によって取得できます。

Python3 の場合

Python3 であれば、bisect モジュールの bisect_left を import した上で、

j = bisect_left(A, b)

によって取得できます。

 

これら標準ライブラリが何をしているのか

C++lower_bound() や、Python3 の bisect_left() は、実は内部で二分探索をやっています。その仕組みに関心のある方は、次の記事などを読んでみてください。

qiita.com

二分探索法の使い方を一通り練習できるサイトとして、アルゴ式で二分探索法のコーナーを作りました。ぜひやってみてください!

algo-method.com

 

コード (C++ と Python3)

以上をまとめて、次のように実装できます。計算量は

  • 数列  A をソートする: O(N \log N)
  • クエリに応える  O(Q \log N)

ですので、全体として  O((N + Q) \log N) となります。

C++

#include <bits/stdc++.h>
using namespace std;

// 無限大を表す値
const long long INF = 1LL << 60;

int main() {
    // 入力
    int N;
    cin >> N;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // ソートしておく
    sort(A.begin(), A.end());

    // 各クエリに答える
    int Q;
    cin >> Q;
    while (Q--) {
        long long B;
        cin >> B;

        // A[j] >= B となる最小の j を求める
        int j = lower_bound(A.begin(), A.end(), B) - A.begin();

        // A[j] と A[j-1] を比較する
        long long res = INF;
        if (j < N) res = min(res, abs(B - A[j]));
        if (j > 0) res = min(res, abs(B - A[j-1]));

        cout << res << endl;
    }
}

Python3

from bisect import bisect_left

# 無限大を表す値
INF = 2 ** 60

# 入力
N = int(input())
A = list(map(int, input().split()))

# ソートしておく
A.sort()

# 各クエリに答える
Q = int(input())
for _ in range(Q):
    B = int(input())

    # A[j] >= B となる最小の j を求める
    j = bisect_left(A, B)

    # A[j] と A[j-1] を比較する
    res = INF
    if j < N:
        res = min(res, abs(B - A[j]))
    if j > 0:
        res = min(res, abs(B - A[j-1]))

    print(res)