D - Kth Excluded Editorial by TumoiYorozu


公式解説では1クエリにつき lower_bound 1回で求めていましたが、
二分探索+upper_bound でも、計算量が \(O(K\log K \log Q)\) と悪化してしまいますが、十分高速に求めることが出来ます。

方針としてはある整数 \(m\) が何番目(以上)のExcludedかは、配列 \(A\) に含まれる \(m\) 以下の個数 \(d\) が求められれば \(m-d\) として求めることができ、この値で二分探索する事ができます。

解答例

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;
using ll = long long;

int main(){
    ll n, q;
    cin >> n >> q;
    vector<ll> a(n);
    for(ll&v:a) cin >> v;
    while(q--) {
        ll k;
        cin >> k;
        ll ok = 2e18;
        ll ng = 0;
        while(ok-ng > 1) {
            ll m = (ok+ng)/2;
            ll d = upper_bound(a.begin(), a.end(), m) - a.begin();
            ll nth = m - d;
            if(nth >= k) {
                ok = m;
            } else {
                ng = m;
            }
        }
        cout << ok << endl;
    }
}

posted:
last update: