【ABC205】D - Kth Excluded

問題リンク

解法

公式解説では2分探索の解法が紹介されていましたが、自分はクエリ先読みで解きました。

$ K_1, K_2, … , K_Q $ と $ A_1, A_2, …, A_N $ をそれぞれ昇順にソートして、先頭から順に突き合わせて、クエリの答えを先に求めておくことを考えます。
クエリの答えを $ Q_1, Q_2, …, Q_i, …, Q_N $ とでもすると、$ K_i < A_1 $ ならば $ Q_i = K_i $ , $ K_i > A_N $ ならば $ Q_i = A_N + N $ というのはすぐにわかります。
後は$ A_1 ≦ K_i ≦ A_N $ のパターンですが、先頭から順にクエリを見ていく時に、今までに配列Aの何番目までを見たか、という変数と$ A_1, A_2, …, A_N $ に含まれない数を何個通りすぎたかという変数を保持することによって、配列Aの要素のどこの間にあるかを特定することができます。

実装

実装はかなり手間取りました。

void solve(long long N, long long Q, std::vector<long long> A, std::vector<long long> K) {
    // クエリの順番を保持するために別の配列にコピーしておく
    vector<ll> q(Q);
    REP (i, Q) q[i] = K[i];

    // ソートする
    ALL(sort, A);
    ALL(sort, K);
    map<ll, ll> k;

    ll j = 0; // 配列Aの添字
    ll pre = A[0] - 1; // 配列Aに含まれない数をいくつ通りすぎたか
    REP (i, Q) {
        ll cnt = K[i] - pre; // 残りいくつの配列Aに含まれない数を通りすぎる必要があるか

        if (cnt <= 0) {
　　　  // 既に通りすぎている場合、適切な位置に戻る
            k[K[i]] = A[j] - (pre - K[i]) - 1;
        } else if (j < N) {
            j++;
            // しゃくとり法みたいな感じ
            while (j < N) {
                ll tmp = A[j] - A[j-1] - 1; // 二つの要素の間にいくつ自然数が含まれるか
                if (cnt - tmp <= 0) {
                    k[K[i]] = A[j-1] + cnt;
                    pre += tmp;
                    break;
                } else {
                    pre += tmp;
                    cnt -= tmp;
                }

                j++;
            }
            if (j >= N) { k[K[i]] = A[N-1] + cnt; }
        } else {
            // 既に配列の全ての要素を飛び越えている場合
            k[K[i]] = A[N-1] + cnt;
        }
    }

    // 出力
    REP (i, Q) {
        c(k[q[i]])
    }
}