問題

公式問題はコチラ。

問題文

長さ $N$ の正整数列 $A = (A_1, A_2, …, A_N)$ と $Q$ 個のクエリが与えられる。$i (1 \leq i \leq Q)$ 番目のクエリでは正整数 $K_i$ が与えられるので、$A_1, A_2, …, A_N$ のいずれとも異なる正整数のうち、小さい方から数えて $K_i$ 番目のものを求めよ。

制約

• $1 \leq N, Q \leq 10^5$
• $1 \leq A_1 < A_2 < … < A_N \leq 10^{18}$
• $1 \leq K_i \leq 10^{18}$
• 入力は全て整数

解説

公式解説はコチラ。

二分探索というのがよくわかっていなくて、TLEしてしまいました。計算量が最大で $O(QN)$ になっていたのが原因でした。

まず、$A_1, A_2, …, A_N$ のいずれでもない正整数を「良い整数」と定義し、$A_i$ 以下の良い整数の個数を $C_i$ とします。（$A_0 = 0, C_0 = 0$）

このとき、$K$ の値によって、以下の2つのパターンが考えられます。

• $C_N < K$ のとき
• $C_N \geq K$ のとき

$C_N < K$ のとき

前者のときは二分探索は必要なく、答えは $A_N + (K \, – \, C_N)$ となります。 $A_N$ 以上の整数はすべて良い整数なので、$A_N$ に $K \, – \, C_N$ を足した数が答えとなります。

たとえば以下の場合を考えます。

$N = 3$
$A = 1 ,2, 4$
$K = 4$

$N = 3$ なので $C_3$ を考えると、$C_3 = 1$ となります。このときの答えは、4 + (4 – 1) で答えは 7 になります。たしかに、良い整数は 3, 5, 6, 7, …で小さい方から数えて 4 番目の値は 7 となっているので、計算は正しそうです。

$C_N \geq K$ のとき

これがお待ちかねの二分探索を使う場合ですね。二分探索によって、$C_i \geq K$ を満たす最小の $C_i$ が求まれば $A_{i-1} + (K \, – \, C_{i-1})$ で先と同様に答えが求まります。ここでは、$A_{i-1} < B < A_{i}$ を満たす整数 $B$ はすべて良い整数という事実を使っています。

C++ では 以下の3つ二分探索用の関数が用意されています。

• binary_search()：探索したい key があるかどうかを返す。
• lower_bound()：探索したい key 以上のイテレータの内、最小のものを返す。
• upper_bound()：探索したい key より大きいイテレータの内、最小のものを返す。

今回の用途では、lower_bound() を使います。

$N = 4$
$A = 1 ,4, 6, 9$
$K = 3$

たとえば上記の例だと、良い整数は 2, 3, 5, 7, 8… で $C_0 = 0, C_1 = 0, C_2 = 2, C_3 = 3, C_4 = 5$ となります。lower_bound() で 3 番目のイテレータが出力されます。$A_2 + (K \ – \ C_2) = 4 + (3 \ – \ 2) \ = \ 5$ となるため、正しい答えが導き出されていることがわかります。

二分探索について

二分探索とは、要素数 $N$ の配列から指定された数を探索するアルゴリズムで、計算量を $O(\log N) $ に抑えられるものです。

wikipedia とか qiita とかが参考になると思います！

ソースコード

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

int main(void) {
    int N, Q;
    vector<long long> A;
    vector<long long> C;
    vector<long long> K;
    cin >> N >> Q;

    A.push_back(0);
    C.push_back(0);
    for (int i=1; i<=N; i++) {
        long long a;
        cin >> a;
        A.push_back(a);
        C.push_back(C[i-1] + A[i] - A[i-1] - 1);
    }
    for (int i=0; i<Q; i++) {
        long long k;
        cin >> k;
        K.push_back(k);
    }

    for (int i=0; i<Q; i++) {
        auto it = lower_bound(C.begin(), C.end(), K[i]);
        if (C[N] < K[i]) cout << A[N] + K[i] - C[N] << endl;
        else cout << A[it-C.begin()-1] + K[i] - C[it-C.begin()-1] << endl;
    }
    
    return 0;
}

感想

「にぶたん」とよく略される二分探索についてまとめてみました。二分探索のアルゴリズムの詳細にはふれず、どちらかというと二分探索を活用したAtCoderの解法についてまとめました。