3261. Count Substrings That Satisfy K-Constraint II

You are given a binary string s and an integer k.

You are also given a 2D integer array queries, where queries[i] = [li, ri].

A binary string satisfies the k-constraint if either of the following conditions holds:

• The number of 0‘s in the string is at most k.
• The number of 1‘s in the string is at most k.

Return an integer array answer, where answer[i] is the number of substrings of s[li..ri] that satisfy the k-constraint.

struct Seg {
    long long mi,ma,sum,lazy;
    Seg *le, *ri;
    Seg(long long l, long long r) : mi(l), ma(r), sum(0), lazy(0), le(nullptr), ri(nullptr) {
        if(l^r) {
            long long m = l + (r - l) / 2;
            le = new Seg(l,m);
            ri = new Seg(m+1,r);
        }
    }
    void propagate() {
        if(lazy) {
            sum += lazy * (ma - mi + 1);
            if(le) le->lazy += lazy;
            if(ri) ri->lazy += lazy;
            lazy = 0;
        }
    }
    void update(long long l, long long r, long long op) {
        propagate();
        if(l <= mi and ma <= r) {
            lazy += op;
            propagate();
            return;
        }
        if(mi > r or ma < l) return;
        le->update(l,r,op);
        ri->update(l,r,op);
        sum = le->sum + ri->sum;
    }
    long long query(long long l, long long r) {
        propagate();
        if(l <= mi and ma <= r) return sum;
        if(mi > r or ma < l) return 0;
        return le->query(l,r) + ri->query(l,r);
    }
};
class Solution {
public:
    vector<long long> countKConstraintSubstrings(string& s, int k, vector<vector<int>>& queries) {
        vector<long long> res(queries.size()), le(s.length());
        Seg* seg = new Seg(0,s.length());
        {
            unordered_map<int,int> seen[2];
            seen[0][-1] = seen[1][-1] = seen[0][0] = seen[1][0] = -1;
            int a = 0, b = 0;
            for(int i = 0; i < s.length(); i++) {
                if(s[i] == '0') {
                    a++;
                    seen[0][a] = max(seen[0][a], i);
                }
                else {
                    b++;
                    seen[1][b] = max(seen[1][b], i);
                }

                int mi = min(seen[0][max(-1,a-k)], seen[1][max(-1,b-k)]) + 1;
                le[i] = mi;
                seg->update(mi,i,1);
            }
        };
        vector<array<long long, 3>> Q;
        for(int i = 0; i < queries.size(); i++) {
            long long l = queries[i][0], r = queries[i][1], now = 0;
            Q.push_back({r,l,i});
        }
        sort(begin(Q), end(Q));
        long long clean = le.size() - 1;
        while(Q.size()) {
            auto [r,l,idx] = Q.back(); Q.pop_back();
            while(clean > r) {
                seg->update(le[clean], clean, -1);
                --clean;
            }
            res[idx] = seg->query(l,r);
        }
        return res;
    }
};