lower_bound を使った解き方【C++】【AtCoder】

C/C++ AtCoder

概要

競技プログラミングにて二分探索について勉強しているときにC++にはlower_boundという関数が用意されていた。
せっかくなので、lower_bound関数の使い方を調べてAtCoderの問題を1問解いてみた。

使い方

公式

cpprefjp.github.io

vector<int> v(n);
int value;
sort(a.begin(), a.end());
auto iter = lower_bound(v.begin(), v.end(), value);

上記のC++日本語リファレンスによるとあらかじめソートされたvectorの中から要素となる value の値以上のイテレータを戻り値として返す。
計算量は二分探索のアルゴリズムなので、 O(logN) と高速。

また、 イテレータをインデックスとして扱いたい場合には

size_t pos = distance(v.begin(), iter);

とdistance関数を使うことでposにどの位置で見つけたか格納することができる。

活用例

競プロ典型90問の内の1問から

atcoder.jp

問題文

ABC 競技プログラミング塾にはN個のクラスがあり、番号 i(1≤i≤N)のクラスはレーティング A_{i}の生徒を対象としています。
今、Q人の生徒がこの塾に入塾しようとしています。
番号j(1≤j≤Q)の生徒のレーティングは B_{j}です。
各生徒は自分に合わないレベルのクラスに行くと不満になります。
各生徒について、その不満度は次のように定義されます。

・対象レーティングaと自分のレーティングbの差の絶対値、すなわち∣a−b∣

j=1,2,3,…,Qそれぞれについて、番号jの生徒の不満度として考えられる最小値を求めてください。
ただし、1人も入らないクラス、複数人から成るクラスがあっても良いものとします。

制約

・1≤N≤300000
・1≤Q≤300000
・0≤A_{i}≤10^{9}
・0≤B_{j}≤10^{9}
・入力はすべて整数

入力

N  
A_{1}　A_{2}　A_{3}　\cdots　A_{N}
Q  
B_{1}  
B_{2}  
B_{3}  
\vdots  
B_{Q}

出力

j=1,2,3,\cdots,Q それぞれについて、標準出力に答えを一行ずつ、総計Q行に出力してください。

入力例

4
4000 4400 5000 3200
3
3312
2992
4229

出力例

112
208
171

解法

上記の問題の解くときの考え方として

  1. B_{i}に対して配列 Aの中で一番近い数字を見つける。

  2. 制約が N、 Qともに300000と大きいため二分探索を使って数字を見つける。

以上、2点を踏まえてプログラムに書き起こすと下記の通り。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int main()
{
    int n, q;
    cin >> n;
    vector<int> a(n);
    for(int i=0; i<n; i++) cin >> a[i];
    sort(a.begin(), a.end());
    cin >> q;
    for(int i=0; i<q; i++){
        int b;
        cin >> b;
        // 二分探索
        auto iter = lower_bound(a.begin(), a.end(), b);
        int idx = distance(a.begin(), iter);
        if(idx == 0) cout << abs(a[idx] - b) << endl;
        else cout << min(abs(a[idx] - b), abs(a[idx-1] - b)) << endl; 
    }
}

計算量も O(NlogN)となるので間に合う。