AtCoder Regular Contest 033 C - データ構造

問題

　数の集合 $S$ に対する以下のクエリを処理せよ。

タイプ1 : $S$ に数 $X$ を追加する。
タイプ2 : $S$ に含まれる数のうち $X$ 番目に小さい数を答え、その数を $S$ から削除する。

解法

　「 $X$ 番目に小さい数⇔それ以下の数が $X$ 個以上追加されている数のうち最も小さい数」と考え、ある数以下の数が $S$ 内にいくつあるかを高速に取得できれば二分探索で答えるべき数を求めることができる。これはセグメント木を使うことにより実装が可能であり、タイプ1である数が登場した際の操作、タイプ2において累積和を求める操作を各 $O(\log{N})$ で行えるためこの問題が解けた。

反省

　ここ数日、あまり机に向かう時間は取れなかったが移動中などにこの問題を考えていて、計数時間は考えたと思うが解けなかったので解説スライドを見た。セグメント木でやっていく発想が一切出てこなかったのは猛省するべき案件で、データ構造という問題名があまりにも露骨だったのでそういう解法を捨ててしまっていた。

　セグメント木の実装は今まで何回かやってきたが、いつもセグメント木をソラで書きたいあなたにを写してばっかりで何も学びを得ていなかったため、今回はかなり自分で考えながら実装してみた。といっても見た覚えがあるのでどうしても似た書き方になってしまうが……。

　数十分かけてなんとか実装できたと思ったが、被覆関係の処理が間違っていて $O(\log{N})$ になっておらずTLE。仕方がなく先のページを見てそこだけ参考にした。

　今までセグメント木を実装できなかったのでセグメント木を使いそうな問題に対してもそうではない解法を考えてしまっていたが、さすがに逃げ続けるわけにはいかないのでちゃんと選択肢に入れて考察していきたい。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = int64_t;

class SegmentTree {
public:
    SegmentTree(ll n) {
        ll radicand = (ll)ceil(log2(n));
        size_ = (ll)pow(2, radicand);
        nodes_.resize(2 * size_ - 1, 0);
    }
    void add(ll x, ll v) {
        for (ll i = x + size_ - 1; ; i = (i - 1) / 2) {
            nodes_[i] += v;
            if (i == 0) {
                break;
            }
        }
    }
    ll getSum(ll a, ll b, ll k = 0, ll l = 0, ll r = -1) {
        //[a, b)の和を求める
        if (r == -1) {
            r = size_;
        }
        assert(a < b);
        assert(l < r);
        //printf("(%lld, %lld, %lld, %lld, %lld)\n", a, b, k, l, r);

        //nodes_[k]が担当する範囲が[l, r)
        if (b <= l || r <= a) {
            //被覆なし
            return 0;
        }

        if (a <= l && r <= b) {
            //完全被覆
            return nodes_[k];
        }

        //部分被覆
        ll kl = (k + 1) * 2 - 1;
        ll kr = (k + 1) * 2;
        return getSum(a, b, kl, l, (l + r) / 2) + getSum(a, b, kr, (l + r) / 2, r);
    }
private:
    ll size_;
    vector<ll> nodes_;
};

int main() {
    ll Q;
    cin >> Q;

    vector<ll> T(Q), X(Q);
    for (ll i = 0; i < Q; i++) {
        cin >> T[i] >> X[i];
    }

    constexpr ll MAX = (ll)2e5;
    SegmentTree st(MAX + 1);
    for (ll i = 0; i < Q; i++) {
        if (T[i] == 1) {
            st.add(X[i], 1);
        } else {
            ll low = 0, high = MAX + 1;
            while (low + 1 != high) {
                ll mid = (low + high) / 2;
                (st.getSum(0, mid) >= X[i] ? high = mid : low = mid);
            }
            cout << low << endl;
            st.add(low, -1);
        }
    }
}