Codeforce's multiset solution using segment tree

multiset.cpp

//
//  multiset.cpp
//  Competitive Programming
//
//  Created by Shahed Bin Satter on 24/9/20.
//  Copyright © 2020 Shahed Bin Satter. All rights reserved.
//

#include <stdio.h>
#include <iostream>

using namespace std;

const int nmax = 1e6 + 1;
int tree[nmax * 4];
int arr[nmax];

void update(int id, int left, int right, int k, int val) {
    if (left == right) {
        tree[id] += val;
//        arr[k] ++;
        return;
    }
    int mid = (left + right) / 2;
    if (k <= mid) {
        update(2 * id, left, mid, k, val);
    } else {
        update(2 * id + 1, mid + 1, right, k, val);
    }
    tree[id] = tree[2 * id] + tree[2 * id + 1];
}

int query(int id, int left, int right, int k) {
    if (left == right) {
        return left;
    }
    int mid = (left + right) / 2;
    if (tree[2 * id] >= k) {
        // left
        return query(2 * id, left, mid, k);
    } else {
        // right
        return query(2 * id + 1, mid + 1, right, k - tree[2 * id]);
    }

}

void build(int id, int left, int right) {
    if (left == right) {
        tree[id] = arr[left];
        return;
    }
    int mid = (left + right) / 2;
    build(2 * id, left, mid);
    build(2 * id + 1, mid + 1, right);
    tree[id] = tree[2 * id] + tree[2 * id + 1];
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    int n, q;
    cin >> n >> q;
    
    for (int i = 1; i <= n; i++) {
        int input;
        cin >> input;
        arr[input] ++;
    }
    build(1, 1, n);
    for (int i = 0; i < q; i++) {
        int k;
        cin >> k;
        if (k > 0) {
            // insert
            update(1, 1, n, k, 1);
        } else {
            // remove
            int elem = query(1, 1, n, -k);
            update(1, 1, n, elem, -1);
        }
    }
    if (tree[1] <= 0) {
        cout << 0 << "\n";
    } else {
        cout << query(1, 1, n, 1) << "\n";
    }
    
    return 0;
}