Sqrt decomposition

When a problem has too many queries or updates to solve with a standard algorithm, sqrt decomposition divides the work into two cases: process large groups as special cases and handle small groups with a direct algorithm.

The idea is straightforward. Sort items by frequency, size, or importance. Items appearing more than sqrt(N) times are "heavy" and receive special treatment. Items appearing less frequently are "light" and are processed directly. This division splits the work: heavy items contribute O(sqrt(N)) per query, and light items contribute O(sqrt(N)) to the total count.

Heavy and light classification

Divide a sequence of N elements into two groups. An element is heavy if it appears more than sqrt(N) times. Otherwise it is light.

There are at most sqrt(N) distinct heavy elements, because if each appears at least sqrt(N) times, the total would exceed N.

The total count of light elements is at most N. But since each light element appears at most sqrt(N) times, the total number of distinct light elements is at most N / sqrt(N) = sqrt(N).

This balance between heavy and light elements is what makes sqrt decomposition work. You cannot have too many heavy elements, and you cannot have too many distinct light elements. So the total work is bounded by sqrt(N) in both directions.

This classification means you can precompute answers for all heavy elements and answer queries about them in O(1). For light elements, you can afford to scan all of them because there are so few distinct ones.

Sqrt blocking on a sequence

Divide a sequence into blocks of size sqrt(N). To answer a range query, you compute the answer for complete blocks in O(sqrt(N)) and handle partial blocks at the boundaries directly.

This is slower than a segment tree but uses less memory and is easier to implement, especially when the operation is not easily composable.

1import math2import sys3
4
5def main() -> None:6    input_data = sys.stdin.read().split()7    idx = 08    n = int(input_data[idx])9    idx += 110    a = [int(input_data[idx + i]) for i in range(n)]11    idx += n12    q = int(input_data[idx])13    idx += 114
15    block_size = int(math.sqrt(n)) + 116    num_blocks = (n + block_size - 1) // block_size17
18    # Precompute block sums19    block_sum = [0] * num_blocks20    for i in range(n):21        block_sum[i // block_size] += a[i]22
23    def range_sum(l: int, r: int) -> int:24        left_block = l // block_size25        right_block = r // block_size26
27        if left_block == right_block:28            # The whole range sits inside one block: no shortcut helps.29            return sum(a[l:r + 1])30
31        # Partial sum in the leftover part of the left block.32        total = sum(a[l:(left_block + 1) * block_size])33        # Whole blocks in the middle, read from the precomputed sums.34        for b in range(left_block + 1, right_block):35            total += block_sum[b]36        # Partial sum in the leftover part of the right block.37        total += sum(a[right_block * block_size:r + 1])38        return total39
40    answers = []41    for _ in range(q):42        l = int(input_data[idx])43        r = int(input_data[idx + 1])44        idx += 245        answers.append(str(range_sum(l, r)))46
47    print("\n".join(answers))48
49
50if __name__ == "__main__":51    main()

8
3 1 4 1 5 9 2 6
3
1 4
0 7
2 2

11
31
4

The program divides the array into blocks and precomputes each block's sum once. Each query then accumulates the sums of the complete blocks it spans, and adds the partial sums from the incomplete blocks at its two edges directly. The first query, 1 to 4, spans part of block 0 and part of block 1 with no complete block between them. The second, 0 to 7, spans one complete block in the middle. The third, 2 to 2, sits entirely inside a single block, so the direct sum is all you need.

Handling updates

When you have both updates and queries, sqrt decomposition becomes more useful. Update a single element in O(1) and recompute its block sum in O(sqrt(N)). A range query still takes O(sqrt(N)).

The cost per operation is roughly the square root of the array size, which is much smaller than N for large inputs but larger than log N. The advantage of sqrt decomposition is that it requires less code and has better cache locality than a segment tree, because you iterate through contiguous memory.

A segment tree is faster in theory (O(log N) per operation) but sqrt decomposition is simpler and can be faster in practice on PyPy due to cache locality and lower constant factors. For problems with N up to 10^5, both are viable. For N = 10^6, sqrt decomposition may TLE.

A second worked example: counting distinct values

The heavy and light split from earlier applies just as well to a harder query. Suppose you have an array of N integers and Q queries, each asking how many distinct values appear in a range.

Scanning the range for each query costs O(N) per query, O(N times Q) overall. Blocking helps the same way it did for sums: divide the array into blocks of size sqrt(N) and precompute the set of distinct values inside each block. A query then merges the precomputed sets of the complete blocks it spans with a direct scan of the two partial blocks at its edges.

Merging sets is more expensive than adding numbers, so the saving is smaller here than it was for range sums, but the shape of the argument does not change: you pay once, up front, for every block, and every later query only touches O(sqrt(N)) of that precomputed work.

Common mistakes

A frequent error is incorrect block boundary handling. If you have N elements and sqrt(N) block size, the last block may have fewer elements. When iterating from one block to another, ensure you handle partial blocks at the start and end of the range. Off-by-one errors here cause wrong answers.

Another mistake is forgetting to update block metadata after updates. If an element changes and belongs to a block, you must recompute the block's aggregate (sum, count, or set). If you forget, subsequent queries will be wrong. Mark the affected block as "dirty" and recompute it lazily if needed.

A third pitfall is using the wrong sqrt threshold. The optimal block size is roughly sqrt(N), but the constant factors vary. For some problems, sqrt(N) is too coarse or too fine. Experiment with slightly different block sizes if your solution is TLE. However, this optimisation should come last, not first.

Why sqrt decomposition works and its cost

Sqrt decomposition works because it balances two competing needs: precomputation cost and query cost. If you precompute too much (e.g., all pairs), you spend O(N^2) time and space upfront. If you precompute too little (e.g., nothing), each query is slow. Sqrt blocking strikes a middle ground: O(N sqrt(N)) preprocessing and O(sqrt(N)) per query.

The cost per operation is O(sqrt(N)), which is slower than logarithmic structures like segment trees (O(log N)) but faster than linear scans (O(N)). For N up to 10^4, sqrt decomposition is competitive. For N = 10^6, the constant factors become important, and PyPy's performance matters. If your solution TLEs on the full problem, try a segment tree or accept partial marks on smaller subtasks.