CP Notebook

← all categories

Sqrt Decomposition 2bd92d13

Sqrt decomposition over a static-size array: range add + range sum via per-block lazy tags and precomputed block sums. Simpler and more flexible than a segment tree when the op is awkward to express as a tree merge (e.g. per-element rebuild rules) — adapt the two inner loops for a different op.

Time: O(\sqrt n) per operation. tested

content/data-structures/sqrt-decomposition.h

template <typename T>
struct SqrtDecomp {
    int n, blockSz;
    vector<T> a, blockSum, lazy;

    SqrtDecomp(vector<T> &arr) : n((int)arr.size()), a(arr) {
        blockSz = (int)sqrt((double)n);
        if (blockSz < 1) blockSz = 1;  // note: max(1, ...) mismatches under #define int long long,
                                        // since the literal 1 stays a real int — see CONVENTIONS.md
        int nb = (n + blockSz - 1) / blockSz;
        blockSum.assign(nb, T());
        lazy.assign(nb, T());
        for (int i = 0; i < n; i++) blockSum[i / blockSz] += a[i];
    }

    void update(int l, int r, T v) {  // inclusive [l, r]
        int bl = l / blockSz, br = r / blockSz;
        if (bl == br) {
            for (int i = l; i <= r; i++) a[i] += v;
            blockSum[bl] += v * (r - l + 1);
            return;
        }
        for (int i = l; i < (bl + 1) * blockSz; i++) a[i] += v;
        blockSum[bl] += v * ((bl + 1) * blockSz - l);
        for (int b = bl + 1; b < br; b++) lazy[b] += v;
        for (int i = br * blockSz; i <= r; i++) a[i] += v;
        blockSum[br] += v * (r - br * blockSz + 1);
    }

    T query(int l, int r) {  // inclusive [l, r]
        int bl = l / blockSz, br = r / blockSz;
        if (bl == br) {
            T res = T();
            for (int i = l; i <= r; i++) res += a[i] + lazy[bl];
            return res;
        }
        T res = T();
        for (int i = l; i < (bl + 1) * blockSz; i++) res += a[i] + lazy[bl];
        for (int b = bl + 1; b < br; b++) res += blockSum[b] + lazy[b] * blockSz;
        for (int i = br * blockSz; i <= r; i++) res += a[i] + lazy[br];
        return res;
    }
};