ABC205 D - Kth Excluded - かんプリンの学習記録

問題概要
解説
提出プログラム
感想

問題概要

長さ ${N}$ の正整数列 ${A}$ が与えられる．以下の ${Q}$ 個のクエリに答えよ．

正整数 ${K}$ が与えられる． ${A}$ のいずれとも異なる正整数のうち，小さい方から数えて ${K}$ 番目のものを求めよ．

${1\leq N,Q\leq 10^5\\ 1\leq A_1< A_2 < \cdots < A_N \leq 10^{18}\\ 0\leq K_i\leq 10^{18}}$

解説

クエリを先読みすることで二分探索を用いずに解きます．

まず， ${K_i}$ を昇順にソートします． ${A_i}$ の昇順に， ${A_i}$ 以下の ${A}$ に含まれない正整数の数 ${S}$ を管理します．例えば， ${A=(3,5,6,7)}$ だと ${S}$ は順に ${2\rightarrow 3\rightarrow 3\rightarrow 3}$ となっていきます．
${K_i\leq S}$ となるとき， ${K_i}$ 番目の正整数はその地点より左側に存在することになるので処理を行うことができます．例として， ${A=(3,5,6,7),K=(2,3)}$ だと， ${S=2}$ となる時点で ${K_1\leq S}$ となっているので ${K=2}$ は ${0}$ と ${A_1}$ の間(両端は含まない)にあることが分かり， ${S=3}$ となる時点では ${K_2\leq S}$ となっているので ${K=3}$ は ${A_1}$ と ${A_2}$ の間にあることが分かります．
実装では ${A_0=0}$ とした方が楽です．詳しくは以下のプログラムを参照してください．計算量は ${O(N+Q\log{Q})}$ となります．

提出プログラム

#include "bits/stdc++.h" 
using namespace std; 
typedef long long ll;

int main() {
    int n,q;cin >> n >> q;
    vector<ll> a(n),k(q),ans(q);
    vector<int> id(q); // id[i]:=昇順i番目のクエリの番号
    for (int i = 0; i < n; i++) cin >> a[i];
    for (int i = 0; i < q; i++) cin >> k[i];
    iota(id.begin(), id.end(),0);
    sort(id.begin(), id.end(),[&](int a,int b){return k[a]<k[b];});
    int t = 0;
    ll sum = 0,pre = 0;// sum:Aにない正整数の数 pre:ひとつ前のA_i
    for (int i = 0; i < n; i++) {
        // K <= S となるKをすべて処理
        while(t < q && k[id[t]] <= sum+a[i]-pre-1) {
            ans[id[t]] = k[id[t]]-sum+pre;
            t++;
        }
        // 値の更新
        sum += a[i]-pre-1;
        pre = a[i];
    }
    for (int i = t; i < q; i++) ans[id[i]] = k[id[i]]-sum+pre;
    for (int i = 0; i < q; i++) cout << ans[i] << endl;
    return 0;
}

問題概要

解説

提出プログラム

感想