kth_segment_tree.cpp

#include<bits/stdc++.h>
using namespace std;
#define int long long

const int N = (int)3e5 + 10;
const int INF = (int)1e18;
int a[N];

class kST{
    int t[4 * N];
    public:
    void build(int node, int l, int r){
        if(l == r){
            t[node] = a[l];
            return;
        }
        int mid = (l + r) / 2;
        build(2 * node, l, mid);
        build(2 * node + 1, mid + 1, r);
        t[node] = t[2 * node] + t[2 * node + 1];
    }
    void update(int node, int l, int r, int i, int val){
        if(l > i || r < i)  return;
        if(l == r){
            t[node] = val;
            return;
        }
        int mid = (l + r) / 2;
        update(2 * node, l, mid, i, val);
        update(2 * node + 1, mid + 1, r, i, val);
        t[node] = t[2 * node] + t[2 * node + 1];
    }
    int sum(int node, int l1, int r1, int l, int r){
        if(l1 > r || r1 < l)    return 0;
        if(l1 >= l && r1 <= r)  return t[node];
        int mid = (l1 + r1) / 2;
        return sum(2 * node, l1, mid, l, r) + sum(2 * node + 1, mid + 1, r1, l, r);
    }
    int kth(int node, int l1, int r1, int l, int r, int k){
        if(l1 > r || r1 < l)    return -1;
        if(l1 == r1)    return l1;
        int mid = (l1 + r1) / 2;
        int left_sum = sum(node * 2, l1, mid, l, r);
        if(k <= left_sum)   return kth(2 * node, l1, mid, l, r, k);
        else return kth(2 * node + 1, mid + 1, r1, l, r, k - left_sum);
    }
};

int32_t main(){
    int n, q; cin >> n >> q;
    for(int i = 1; i <= n; i++) cin >> a[i];
    kST seggs;
    seggs.build(1, 1, n);
    while(q--){
        int type; cin >> type;
        if(type == 1){
            int i; cin >> i; i++;
            a[i] = 1 - a[i];
            seggs.update(1, 1, n, i, a[i]);
        }else{
            int k;
            cin >> k; k++;
            int ans = seggs.kth(1, 1, n, 1, n, k);
            cout << ans - 1 << "\n";
        }
    }
    return 0;
}