Educational Codeforces Round 87 (Rated for Div. 2) D - Multiset

コンテスト参加記はこちら。

coonevo.hatenablog.com

D - Multiset

問題へのリンク

問題概要

集合の初期状態が与えられる。 以下2つのクエリを処理したあと、残っている要素のいずれかを出力せよ（空の場合は-1）。

• 要素を追加。

• 番目の要素を削除。

解説

BITを使うと重複のある集合の挿入、削除、番目の要素の取得がそれぞれでできる。

番目の要素の取得は、BITの和を求める操作を使うとになるが、BIT上の二分探索を用いればでできる。

template<typename T>
class BIT{
  public:
    int N;
    vector<T> data;
    BIT(T _N):N(_N){
        data.assign(N+1, 0);
    };
    
    // a is 1-indexed
    void add(int a, T w){
        for(int x = a; x <= N; x += x & -x)data[x] += w;
    }
    // 1-indexed sum of prefix [0, a]
    T sum(int a){
        T res = 0;
        for(int x = a; x > 0; x -= x & -x)res += data[x];
        return res;
    }
    // 1-indexed sum of range [l, r]
    T sum(int l, int r){return sum(r) - sum(l-1);}

    // 0-indexed add
    void add0(int a, T w){add(a + 1, w);}
    // 0-indexed sum
    T sum0(int a){return sum(a + 1);}
    // 0-indexed sum of range
    T sum0(int l, int r){return sum0(r) - sum0(l-1);}
    // k-th number (k is 1 - indexed)
    T get(int k){
        T res = 0;
        int sz = 1;
        while(sz < (int)data.size()) sz <<= 1;
        for(int i = sz / 2; i > 0; i >>= 1){
            if(res + i <= N && data[res + i] < k){
                k -= data[res + i];
                res += i;
            }
        }
        return res + 1;
    }
};

int main() {

    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    vector<int> k(Q);
    rep(i,N)cin >> a[i];
    rep(i,Q)cin >> k[i];
    BIT<ll> B(N);
    rep(i,N)B.add(a[i],1);
    rep(i,Q){
        // 削除クエリ
        if(k[i] < 0){
            k[i] = -k[i];
            B.add(B.get(k[i]), -1);
        }
        // addクエリ
        else{
            B.add(k[i],1);
        }
    }
    rep1(i,N+1){
        if(B.sum(i,i) > 0){
            cout << i << endl;
            return 0;
        }
    }
    cout << 0 << endl;
}