Square Root Decomposition: Range Queries in O(sqrt n)

Sanjeev SharmaSanjeev Sharma
3 min read

Advertisement

Square Root Decomposition

Divide array into blocks of size ~sqrt(n). Precompute block aggregates.

Query [l, r]: O(sqrt n) — partial first block + full middle blocks + partial last block Update i: O(1) — update element + update block aggregate

Python Implementation

import math

class SqrtDecomp:
    def __init__(self, arr):
        self.n = len(arr)
        self.arr = arr[:]
        self.block = int(math.sqrt(self.n)) + 1
        self.num_blocks = (self.n + self.block - 1) // self.block
        self.blocks = [0] * self.num_blocks
        for i, v in enumerate(arr):
            self.blocks[i // self.block] += v

    def update(self, i, val):
        self.blocks[i // self.block] += val - self.arr[i]
        self.arr[i] = val

    def query(self, l, r):  # inclusive [l, r]
        total = 0
        bl, br = l // self.block, r // self.block
        if bl == br:
            return sum(self.arr[l:r+1])
        total += sum(self.arr[l:(bl+1)*self.block])
        for b in range(bl+1, br):
            total += self.blocks[b]
        total += sum(self.arr[br*self.block:r+1])
        return total

# Range minimum query variant
class SqrtMin:
    def __init__(self, arr):
        self.n = len(arr)
        self.arr = arr[:]
        self.B = int(self.n**0.5) + 1
        self.blocks = [float('inf')] * ((self.n + self.B - 1) // self.B)
        for i, v in enumerate(arr):
            self.blocks[i // self.B] = min(self.blocks[i // self.B], v)

    def update(self, i, val):
        self.arr[i] = val
        b = i // self.B
        self.blocks[b] = min(self.arr[b*self.B:(b+1)*self.B])

    def query(self, l, r):
        bl, br = l // self.B, r // self.B
        if bl == br:
            return min(self.arr[l:r+1])
        res = min(self.arr[l:(bl+1)*self.B])
        for b in range(bl+1, br):
            res = min(res, self.blocks[b])
        res = min(res, min(self.arr[br*self.B:r+1]))
        return res

C++ Implementation

#include <bits/stdc++.h>
using namespace std;

struct SqrtDecomp {
    int n, B;
    vector<long long> arr, blocks;

    SqrtDecomp(vector<int>& a) : n(a.size()), B(sqrt(a.size())+1),
        arr(a.begin(), a.end()), blocks((a.size()+sqrt(a.size()))/sqrt(a.size())+1, 0) {
        for (int i = 0; i < n; i++) blocks[i/B] += arr[i];
    }

    void update(int i, long long val) {
        blocks[i/B] += val - arr[i];
        arr[i] = val;
    }

    long long query(int l, int r) {
        long long res = 0;
        int bl = l/B, br = r/B;
        if (bl == br) {
            for (int i = l; i <= r; i++) res += arr[i];
            return res;
        }
        for (int i = l; i < (bl+1)*B; i++) res += arr[i];
        for (int b = bl+1; b < br; b++) res += blocks[b];
        for (int i = br*B; i <= r; i++) res += arr[i];
        return res;
    }
};

Mo's Algorithm (Offline Range Queries)

Sort queries by block of left endpoint, then right endpoint:

def mo_algorithm(arr, queries):
    n = len(arr)
    B = int(n**0.5) + 1
    # Sort: block of l, then r (alternating direction for optimization)
    queries.sort(key=lambda q: (q[0]//B, q[1] if (q[0]//B)%2==0 else -q[1]))

    cur_l, cur_r = 0, -1
    cur_sum = 0
    results = [0] * len(queries)

    for idx, (l, r) in enumerate(queries):
        while cur_r < r: cur_r += 1; cur_sum += arr[cur_r]
        while cur_l > l: cur_l -= 1; cur_sum += arr[cur_l]
        while cur_r > r: cur_sum -= arr[cur_r]; cur_r -= 1
        while cur_l < l: cur_sum -= arr[cur_l]; cur_l += 1
        results[idx] = cur_sum
    return results

Complexity

  • Query/Update: O(sqrt n)
  • Mo's algorithm: O((n + q) * sqrt n) for offline queries

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro