POI Task Meteors

#include <iostream>
#include <vector>
#include <tuple>
#include <algorithm>

using namespace std;

const long long INF = 2e9;

void update(vector<long long> &fenwick, int index, long long value) {
    while (index < (int)fenwick.size()) {
        fenwick[index] += value;
        index += index & -index;
    }
}
void updateRange(vector<long long> &fenwick, int l, int r, long long value) {
    update(fenwick, l, value);
    if (r+1 < (int)fenwick.size()) update(fenwick, r+1, -value);
}
long long query(vector<long long> &fenwick, int index) {
    long long sum = 0;
    while (index) {
        sum += fenwick[index];
        index -= index & -index;
    }
    return sum;
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int n, m; cin >> n >> m;
    vector<vector<int>> stateIndices(n);
    for (int i = 1; i <= m; ++i) {
        int state; cin >> state;
        stateIndices[state-1].push_back(i);
    }

    vector<int> minMeteors(n);
    for (int &x : minMeteors) cin >> x;

    int k; cin >> k;
    vector<tuple<int, int, int>> meteors(k);
    for (auto &met : meteors) cin >> get<0>(met) >> get<1>(met) >> get<2>(met);

    // Parallel Binary Search on number of meteors for range
    vector<int> stateL(n, 0), stateR(n, k+1);
    vector<pair<int, int>> q;
    while (true) {
        bool going = false; // at least one binary search still active
        q.clear();
        for (int s = 0; s < n; ++s) {
            if (stateL[s]+1 < stateR[s]) {
                q.emplace_back((stateL[s]+stateR[s])/2, s);
                going = true;
            }
        }
        sort(q.begin(), q.end());

        if (!going) break;

        int nextQ = 0; // index in q

        // Build segment tree and ask queries along the way
        vector<long long> fenwick(1+m);
        for (int i = 0; i < k; ++i) {
            // Update the segment tree
            int qL = get<0>(meteors[i]), qR = get<1>(meteors[i]), a = get<2>(meteors[i]);
            if (qL <= qR) updateRange(fenwick, qL, qR, a);
            else updateRange(fenwick, qL, m, a), updateRange(fenwick, 1, qR, a);

            // Make any state queries necessary
            while (nextQ < (int)q.size() && q[nextQ].first == i+1) {
                int s = q[nextQ].second;
                long long total = 0;
                for (int index : stateIndices[s]) {
                    total += query(fenwick, index);
                    if (total > INF) break; // to avoid overflow
                }
                if (total < minMeteors[s]) stateL[s] = i+1;
                else stateR[s] = i+1;
                nextQ++;
            }
        }
    }

    for (int s = 0; s < n; ++s) {
        if (stateR[s] == k+1) cout << "NIE\n";
        else cout << stateR[s] << '\n';
    }
    return 0;
}