色んな方法があるが、できるだけ楽にやりたい！

問題へのリンク

問題概要

単調増加な $N$ 要素からなる数列 $A_{1}, A_{2}, \dots, A_{N}$ がある。この数列について、次の $Q$ 回のクエリに答えよ。

【クエリ】
整数 $K$ が与えられるので、 $A_{1}, A_{2}, \dots, A_{N}$ のどれとも異なる数値 (よい数と呼ぶこととする) のうち、小さい方から $K$ 番目の値を求めよ。

制約

$1 \le N, Q \le 10^{5}$

考えたこと

いろんな方法が考えられる。

自分が最初にやったのは「 $x$ 以下のよい数は何個あるか」という判定問題を解くことで、二分探索するという方法だった。この判定問題を解くこと自体に、通常 lower_bound() を活用する方法が考えられるので、全体として $O(N + Q (\log N)^{2})$ の計算量となる。

ここでは、もっと計算量の小さい方法を考える。

どの区間に属するかを求める

この問題では、実は「 $K$ 番目のよい数」が

$A_{1}$ より小さい
$A_{1}$ より大きくて $A_{2}$ より小さい
$A_{2}$ より大きくて $A_{3}$ より小さい
...
$A_{N}$ より大きい

という $N+1$ 種類の区間のうちのどこに属するのかが、比較的簡単に分かるのだ。具体的には、次の配列を定義しよう。

$C_{i}$ ← $A_{i}$ より小さい「よい数」の個数

たとえば、 $A = (3, 5, 6, 10, 20)$ のとき、

$C = (2, 3, 3, 6, 15)$

と求められる。この配列 num において、 $K$ がどこに属するかを考えればよい。たとえば、上の例で $K = 8$ ならば、その値は配列 $C$ において $C_{4} = 6$ と $C_{5} = 15$ の間にあるので、「 $K$ 番目のよい数」は $A_{4} = 10$ と $A_{5} = 20$ の間にあることがわかるのだ。

より正確には、次のようにすればよい。

が以上未満であるとする
- 「 $C_{1}$ 未満」と「 $C_{N}$ 以上」の場合はよしなにする
- これは lower_bound() を用いて求められる
このとき、答えは次のようになる。
- $K$ が $C_{1}$ 未満の場合は、 $K$
- そうでない場合は、 $A_{i} + (K - C_{i})$

この方法を用いれば、全体の計算量は $O(N + Q\log N)$ となる。

コード

全体的に 0-indexed で実装した。

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

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<long long> A(N), C(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        C[i] = A[i] - i - 1;
    }
    
    while (Q--) {
        long long K;
        cin >> K;
        int i = lower_bound(C.begin(), C.end(), K) - C.begin() - 1;
        
        if (i < 0) cout << K << endl;
        else cout << A[i] + (K - C[i]) << endl;
    }
}

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

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

AtCoder ABC 205 D - Kth Excluded (2Q, 茶色, 400 点)

問題概要

制約

考えたこと

どの区間に属するかを求める

コード