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

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

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

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

問題概要

単調増加な  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 の間にあることがわかるのだ。

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


  •  K C_{i} 以上  C_{i+1} 未満であるとする
    •  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;
    }
}