Square Root Decomposition

This is an extension topic beyond most interviews; it pays off in hard OA problems where a segment tree would need an awkward merge or lazy tag, and roughly O(sqrt n) work per operation is fast enough.

Intuition

Square-root decomposition is less a data structure than a way of thinking: cut the array into blocks of about sqrt(n) elements and keep a small summary for each block. Any range [l, r) then consists of three parts:

1. a partial block on the left, handled element by element (at most B elements),
2. some number of whole blocks in the middle, handled through their summaries (at most n / B blocks),
3. a partial block on the right, handled element by element.

With block size B, an operation costs O(B + n / B). Balancing the two terms gives B = sqrt(n) and O(sqrt n) per operation. For n = 10^5 that is a few hundred steps, slower than a segment tree's O(log n) but often far easier to get right.

The summary depends on the question:

• Range add + range sum: each block keeps its sum and a lazy value "every element here has an extra +v". Whole blocks just bump the lazy value; partial blocks update elements and the block sum.
• Range add + "how many elements are >= x in a range": each block keeps a sorted copy of its elements. Adding the same value to a whole block does not change the order inside it, so whole blocks only bump a lazy value, and a query binary searches each whole block for x - lazy. Partial blocks are modified directly and re-sorted. A segment tree has no simple mergeable summary for this query, which is exactly when blocks shine.

The same idea applies to queries instead of elements: buffer pending updates, let each query scan the buffer, and rebuild a static structure (like prefix sums) after every sqrt(q) updates. And to values: treat "heavy" items (frequency above sqrt(n)) differently from "light" ones. Offline query reordering (Mo's algorithm) builds on the same block idea but is beyond interview scope.

When to use it in interviews / OAs

• Range updates plus range queries where the per-range information does not merge nicely or the lazy tags do not compose: counting elements above a threshold under range add, k-th smallest in a range with updates (binary search on the answer on top of counting).
• Constraints around n, q ≤ 10^5 where O((n + q) * sqrt n) fits the time limit (be careful in Python; see below).
• A safer fallback when you are not sure how to design a segment tree's node summary or tag under time pressure.
• Batching: many cheap updates and occasional queries (or the other way around), where periodic rebuilding wins.
• Heavy/light splits: graph or counting problems where a few elements have huge degree or frequency.

Template

Range add and range sum. The left and right partial blocks are walked element by element; whole blocks only touch their lazy value. The last block may be shorter than B, and the loops automatically treat it as a partial block.

import math


class SqrtArray:
    # Range add and range sum in O(sqrt n) per operation.
    def __init__(self, a: list[int]) -> None:
        self.a = list(a)
        self.B = max(1, math.isqrt(len(a)))  # Block size about sqrt(n).
        nb = (len(a) + self.B - 1) // self.B
        self.sum = [0] * nb  # Sum of the stored values in each block.
        self.lazy = [0] * nb  # Pending addition for every element of a block.
        for i, x in enumerate(a):
            self.sum[i // self.B] += x

    def add(self, l: int, r: int, v: int) -> None:  # a[l:r] += v
        B = self.B
        while l < r and l % B:  # Left partial block: element by element.
            self.a[l] += v
            self.sum[l // B] += v
            l += 1
        while l + B <= r:  # Whole blocks: O(1) each.
            self.lazy[l // B] += v
            l += B
        while l < r:  # Right partial block.
            self.a[l] += v
            self.sum[l // B] += v
            l += 1

    def query(self, l: int, r: int) -> int:  # Sum of a[l:r].
        B, s = self.B, 0
        while l < r and l % B:
            s += self.a[l] + self.lazy[l // B]
            l += 1
        while l + B <= r:
            b = l // B
            s += self.sum[b] + self.lazy[b] * B
            l += B
        while l < r:
            s += self.a[l] + self.lazy[l // B]
            l += 1
        return s


sa = SqrtArray([1, 2, 3, 4, 5, 6, 7])
sa.add(1, 6, 10)  # a = [1, 12, 13, 14, 15, 16, 7]
assert sa.query(0, 7) == 78 and sa.query(2, 4) == 27

Range add and "count elements >= x in a range": each block keeps a sorted copy, whole blocks are binary searched, partial blocks are rebuilt.

import bisect
import math


class BlockCounter:
    # Range add, and "how many a[i] >= x for i in [l, r)", with a sorted copy per block.
    def __init__(self, a: list[int]) -> None:
        self.a = list(a)
        self.B = max(1, math.isqrt(len(a)))
        nb = (len(a) + self.B - 1) // self.B
        self.lazy = [0] * nb
        self.sorted = [sorted(self.a[b * self.B : (b + 1) * self.B]) for b in range(nb)]

    def _touch(self, l: int, r: int, v: int) -> None:  # Partial update inside one block, then re-sort it.
        for i in range(l, r):
            self.a[i] += v
        b = l // self.B
        self.sorted[b] = sorted(self.a[b * self.B : (b + 1) * self.B])

    def add(self, l: int, r: int, v: int) -> None:  # a[l:r] += v, requires l < r
        B = self.B
        lb, rb = l // B, (r - 1) // B
        if lb == rb:
            self._touch(l, r, v)
            return
        self._touch(l, (lb + 1) * B, v)
        for b in range(lb + 1, rb):
            self.lazy[b] += v  # Sorted order inside a block survives a uniform shift.
        self._touch(rb * B, r, v)

    def count_at_least(self, l: int, r: int, x: int) -> int:  # requires l < r
        B = self.B
        lb, rb = l // B, (r - 1) // B
        if lb == rb:
            return sum(1 for i in range(l, r) if self.a[i] + self.lazy[lb] >= x)
        cnt = sum(1 for i in range(l, (lb + 1) * B) if self.a[i] + self.lazy[lb] >= x)
        cnt += sum(1 for i in range(rb * B, r) if self.a[i] + self.lazy[rb] >= x)
        for b in range(lb + 1, rb):
            s = self.sorted[b]
            cnt += len(s) - bisect.bisect_left(s, x - self.lazy[b])  # Binary search in the block.
        return cnt


bc = BlockCounter([5, 1, 4, 2, 3, 6, 0, 7, 8])
assert bc.count_at_least(0, 9, 5) == 4
bc.add(1, 7, 3)  # a = [5, 4, 7, 5, 6, 9, 3, 7, 8]
assert bc.count_at_least(0, 9, 5) == 7 and bc.count_at_least(2, 5, 6) == 2

Choosing the block size: the right B balances the two sides of the cost. Here a partial update re-sorts a block (O(B log B)) and a query binary searches every whole block (O((n / B) log B)), so plain sqrt(n) is fine. When one side is more expensive, shift B toward the cheaper side. In Python, larger blocks also let built-ins such as sum(a[i:j]) or list slicing do the per-element work in C.

Complexity

Structure	Update	Query	Space
range add + range sum	O(sqrt n)	O(sqrt n)	O(n)
range add + count >= x	O(sqrt n log n)	O(sqrt n log n)	O(n)
query batching (rebuild every T updates)	O(1) per update, plus an O(n) rebuild every T updates	O(T)	O(n)

For comparison, a segment tree does range add + range sum in O(log n). With n = q = 10^5, q operations cost up to about q * (2B + n / B), roughly 10^7 to 10^8 elementary steps: fine in C++ or Java, too slow in pure Python unless the inner loops use built-ins or the block size is tuned.

Common pitfalls

• Block boundaries. Block b covers [b * B, min((b + 1) * B, n)). The last block can be shorter; never multiply its lazy value by B unless it is full.
• Forgetting the lazy value. Whenever you read or modify individual elements of a block, add the block's pending lazy value (or push it into the elements first and reset it).
• Stale derived data. After a partial update, rebuild that block's summary (sum, sorted copy, counts) before the next query.
• Edge cases. Empty ranges, ranges inside a single block, and n smaller than B.
• Wrong block size. A badly chosen B (say 10 for n = 10^5) silently turns O(sqrt n) into O(n). Derive it from the cost formula.
• Using it when a simpler tool exists. Static data needs prefix sums or a sparse table; sums with point updates need a Fenwick tree.

Interview variations

• Range add, range sum as an alternative to a lazy segment tree.
• Range add, count of elements >= x (or <= x) in a range.
• k-th smallest in a range with updates: binary search on the value and count with the blocks, about O(sqrt n * log n * log V) per query.
• Periodic rebuild: keep a small list of recent updates and rebuild prefix sums or a sorted array every sqrt(q) updates.
• Heavy/light thresholding: handle the few high-frequency values (or high-degree vertices) separately, and brute force the many light ones.