POI Task Meteors (49pts - TLE)

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

using namespace std;

struct Node {
    long long sum = 0, prop = 0;
    Node * left, * right;
};
void propagate(Node * &base) {
    if (!base->left) base->left = new Node();
    if (!base->right) base->right = new Node();
    base->left->sum += base->prop;
    base->right->sum += base->prop;
    base->left->prop += base->prop;
    base->right->prop += base->prop;
    base->prop = 0;
}
void update(Node * &base, long long val, int qL, int qR, int l, int r) {
    if (qR <= l || r <= qL) return;
    if (qL <= l && r <= qR) {
        base->sum += (r-l)*val;
        base->prop += val;
        return;
    }

    propagate(base);
    int m = (l+r)/2;
    update(base->left, val, qL, qR, l, m);
    update(base->right, val, qL, qR, m, r);

    base->sum = base->left->sum + base->right->sum;
}
long long query(Node * &base, int index, int l, int r) {
    if (r-l==1) return base->sum;
    propagate(base);
    int m = (l+r)/2;
    if (index < m) return query(base->left, index, l, m);
    else return query(base->right, index, m, r);
}
void deleteTree(Node * &base) {
    if (base) {
        deleteTree(base->left);
        deleteTree(base->right);
        delete base;
    }
}

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 = 0; 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);
        --get<0>(met), --get<1>(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
        Node * root = new Node();
        auto updateTree = [&root, &m](int l, int r, int a) {
            update(root, a, l, r+1, 0, 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) updateTree(qL, qR, a);
            else updateTree(qL, m-1, a), updateTree(0, 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(root, index, 0, m);
                }
                if (total < minMeteors[s]) stateL[s] = i+1;
                else stateR[s] = i+1;
                nextQ++;
            }
        }
        deleteTree(root);
    }

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