Lab 12 — Sqrt Decomposition & Mo’s Algorithm

Goal

Answer offline range queries in O((n + q)·√n) when the aggregate does not fit a segment tree — the canonical case: “number of distinct values in a[l..r]”. Warm-up: block decomposition for point-update/range-sum. Main event: Mo’s algorithm — sort queries into block order, slide a window with O(1) add/remove. Target: Mo’s from cold start in <20 minutes, movement bound derived on a whiteboard.

Background Concepts

Block decomposition (warm-up). Split the array into blocks of size B ≈ √n; keep one aggregate per block. Point update touches one block: O(1). Range query = two partial blocks scanned element-wise + whole blocks read from the per-block table: O(B + n/B), minimized at B = √n → O(√n). It is a segment tree with exactly two levels — worse asymptotics, trivial code, and it tolerates any per-block precomputation (sorted copies for “count < x in range”, etc.).

Mo’s algorithm. Some aggregates are easy to maintain incrementally but impossible to merge: distinct-count has no combine(left_half, right_half). Mo’s exploits incrementality: maintain a current window [cl, cr] with add(i)/remove(i) in O(1), and answer queries offline in an order that keeps total pointer travel small. Sort queries by (block of l, then r — with r’s direction alternating per block, “snake order”). Movement bound: within one l-block, r is monotone → r travels ≤ n per block, n/B blocks → O(n²/B); l wiggles ≤ B per query plus B per block change → O(qB + n). At B = √n both terms are O((n+q)√n); at B = n/√q the total tightens to O(n√q).

Interview Context

This is a competitive-programming staple that leaks into interviews at HFT shops and infra teams as “design a batch analytics query engine”. The signal being tested: recognizing that distinct-count is not decomposable (a segment tree node cannot summarize it), and that “offline” is a license to reorder work. Saying “Mo’s algorithm, O((n+q)√n), needs all queries upfront” plus the sort trick is the entire expected answer; the code is secondary. Contrast question interviewers love: LC 1157 (Online Majority Element) looks adjacent but is online — different tool entirely (segment tree over Boyer-Moore candidates + binary-searched occurrence lists).

Problem Statement

Given an array a of n integers and q queries (l, r), report the number of distinct values in a[l..r] for each. Canonical judge: SPOJ DQUERY (n ≤ 3×10⁴, q ≤ 2×10⁵). There is no exact LeetCode equivalent — LeetCode’s range-query problems are online or decomposable; say so honestly rather than forcing a bad analogue.

Constraints

  • n ≤ 3×10⁴, q ≤ 2×10⁵ (DQUERY); the technique scales to n, q ≈ 10⁵–10⁶ in compiled languages.
  • Values up to 10⁶ — coordinate-compress so cnt is an array.
  • Queries known upfront (offline); no updates interleaved.
  • Ranges inclusive; DQUERY is 1-indexed (convert at the boundary).

Clarifying Questions

  1. Online or offline? (Offline — all queries available before answering. Mo’s is unusable online.)
  2. Updates between queries? (No. With updates → Mo with time dimension, O(n^{5/3}) — see follow-ups.)
  3. Value range? (Up to 10⁶ — compress to [0, n) so the count table is a flat array.)
  4. Inclusive ranges? 0- or 1-indexed? (Inclusive; this lab is 0-indexed, DQUERY input is 1-indexed.)
  5. Answer order? (Answers must be emitted in the original query order — Mo reorders internally only.)

Examples

a = [1, 1, 2, 1, 3]
query(0, 4) → 3        (values {1, 2, 3})
query(1, 3) → 2        (values {1, 2})          ← DQUERY sample, 0-indexed

Warm-up: a = [3, 1, 4, 1, 5],  sum(1, 3) → 6,  update(2, 0),  sum(1, 3) → 2

Initial Brute Force

Per query, scan the range and build a set; answer is its size.

Brute Force Complexity

O(n) per query → O(n·q) = 6×10⁹ element visits at DQUERY limits, each with a hash insert. TLE by two orders of magnitude. Space O(n) per query for the set.

Optimization Path

Brute force → try a segment tree and fail: distinct-count does not merge (knowing each half has 3 distinct values says nothing about the union) → Mo’s O((n+q)√n) with O(1) add/remove — the general tool whenever the window aggregate is incremental → for distinct-count specifically, an offline BIT sorted by r (“keep a 1 only at each value’s last occurrence ≤ r”) achieves O((n+q) log n) and beats Mo — know it exists; Mo remains the answer for non-BIT-able aggregates (mode, count of pairs, XOR-frequency queries) → online requirement → persistent segment tree (see follow-ups).

Final Expected Approach

from math import isqrt

class SqrtBlocks:
    """Warm-up: point update, range sum. O(1) update, O(√n) query."""

    def __init__(self, a):
        self.a = list(a)
        self.b = isqrt(len(a)) + 1
        self.block = [0] * ((len(a) + self.b - 1) // self.b)
        for i, x in enumerate(a):
            self.block[i // self.b] += x

    def update(self, i, x):                      # a[i] = x
        self.block[i // self.b] += x - self.a[i]
        self.a[i] = x

    def query(self, l, r):                       # sum of a[l..r], inclusive
        s = 0
        while l <= r and l % self.b:             # left partial block
            s += self.a[l]; l += 1
        while l + self.b - 1 <= r:               # whole blocks
            s += self.block[l // self.b]; l += self.b
        while l <= r:                            # right partial block
            s += self.a[l]; l += 1
        return s

def mo_distinct(a, queries):
    """queries: (l, r) inclusive, 0-indexed. Offline. O((n + q)·√n)."""
    n, q = len(a), len(queries)
    if q == 0:
        return []
    B = max(1, int(n / q ** 0.5))                # block size ≈ n / √q
    def key(i):
        l, r = queries[i]
        blk = l // B
        return (blk, r if blk % 2 == 0 else -r)  # snake order: r alternates
    order = sorted(range(q), key=key)
    comp = {v: i for i, v in enumerate(sorted(set(a)))}
    va = [comp[x] for x in a]                    # coordinate-compress values
    cnt = [0] * len(comp)
    distinct = 0
    ans = [0] * q
    cl, cr = 0, -1                               # current window [cl, cr], empty
    for qi in order:
        l, r = queries[qi]
        while cl > l:                            # expand before shrink
            cl -= 1
            cnt[va[cl]] += 1
            if cnt[va[cl]] == 1: distinct += 1
        while cr < r:
            cr += 1
            cnt[va[cr]] += 1
            if cnt[va[cr]] == 1: distinct += 1
        while cl < l:
            cnt[va[cl]] -= 1
            if cnt[va[cl]] == 0: distinct -= 1
            cl += 1
        while cr > r:
            cnt[va[cr]] -= 1
            if cnt[va[cr]] == 0: distinct -= 1
            cr -= 1
        ans[qi] = distinct
    return ans

Data Structures Used

  • block — per-block aggregate array (warm-up), size ⌈n/B⌉.
  • cnt — multiplicity table over compressed values; distinct — one running int. The whole “state” of Mo’s is these two.
  • order — the permutation of query indices; answers written back by original index.

Correctness Argument

Blocks: block[j] = Σ of its elements is maintained exactly by each update (delta applied to one block); a query partitions [l, r] into ≤ 2 partial blocks (scanned directly) and whole blocks (read from block) — a disjoint exact cover, so the sum is exact.

Mo’s: invariant — after the four while-loops, cnt[x] is the multiplicity of x in a[cl..cr] and distinct = |{x : cnt[x] > 0}|, with (cl, cr) = (l, r). Each add/remove changes one multiplicity by 1 and adjusts distinct exactly at the 0↔1 boundary, so the invariant is preserved per step; the loops terminate with the window equal to the query. Correctness is independent of the query order — the sort affects only running time. Expanding (cl↓, cr↑) before shrinking guarantees the window never passes through an inverted state, so no count goes negative.

Movement bound: stated in Background; the key observation is that r resets at most once per l-block (not per query), and snake order removes even those resets.

Complexity

OperationTimeSpace
SqrtBlocks update / queryO(1) / O(√n)O(n)
Mo’s sortO(q log q)O(q)
Mo’s pointer movementO((n + q)·√n) total (O(n√q) with B = n/√q)O(n)

Implementation Requirements

  • cnt must be a flat array over compressed values — a dict is correct but ~5× slower in Python and often the difference between AC and TLE.
  • Block size clamped: max(1, int(n / √q)) (q ≫ n² would otherwise give B = 0 → division by zero in the sort key).
  • Answers stored by original query index; never emit in sorted order.
  • Pointer discipline: exactly the four loops, expand-before-shrink, window inclusive [cl, cr] initialized to (0, −1).
  • Warm-up struct: no lazy anything — that is the point of the exercise; it is 20 lines.

Tests

  • DQUERY sample: a = [1,1,2,1,3], queries (0,4) → 3, (1,3) → 2.
  • Single element, empty query list, all-equal array (every answer 1), all-distinct array (answer = r−l+1).
  • Degenerate blocks: q ≫ n² (B clamps to 1); q = 1 (B = n).
  • SqrtBlocks: interleaved random updates/queries vs sum(a[l:r+1]), including negative values.
  • Stress: 200 random arrays (n ≤ 80, values ≤ 12), all Mo answers vs len(set(a[l:r+1])). (All of the above verified against brute force before this lab shipped.)
  • Perf: n = 3×10⁴, q = 2×10⁵ random queries — count total pointer steps, confirm ≈ (n+q)√n, not n·q.

Follow-up Questions

  • Why does snake order (alternating r direction) help? → With plain (block, r ascending) sorting, r crashes from ~n back to ~1 at every block boundary — n/B resets ≈ √q · n extra steps. Alternating r’s direction per block turns the reset into a continuation; r sweeps up-then-down like a snake. Same asymptotics, roughly 2× fewer moves in practice — often the AC/TLE line.
  • Queries interleaved with point updates? → Mo with a time dimension: each query is (l, r, t) where t = number of updates before it; sort by (l-block, r-block, t) with block size n^{2/3}; three pointers (cl, cr, time), where moving time applies/rolls back one update. Total O(n^{5/3}) for n ≈ q. Painful but mechanical.
  • When does a segment tree or BIT win? → Whenever queries are online, updates are frequent, or the aggregate is decomposable (sum, min, max, gcd — anything with a merge). Mo’s is strictly a last resort for incremental-but-not-mergeable aggregates on offline queries. For distinct-count itself, the offline sort-by-r BIT trick (a 1 at each value’s last occurrence) is O((n+q) log n) — faster than Mo; Mo earns its keep on mode, “count pairs with equal values”, and friends.
  • Hilbert curve ordering? → Map each query (l, r) to its position d(l, r) on a Hilbert curve over the n×n grid and sort by d alone. Locality in d implies locality in both coordinates simultaneously — provably O(n√q) movement without tuning B, and measurably faster than block sort at large q. Standard drop-in in CP templates.
  • Same query, but online? → Persistent segment tree: version r of the tree holds a 1 at the last occurrence ≤ r of each value; distinct(l, r) = range-sum of [l, r] in version r. O(log n) per query, O(n log n) space, fully online — persistent structures are covered in phase-12. This is the standard escape hatch when the interviewer removes the “offline” concession.
  • Why B = n/√q instead of √n? → Balancing the two movement terms n²/B (r-travel) and qB (l-travel) gives B = n/√q, total O(n√q). At q ≈ n they coincide at √n; at q ≫ n the tuned block is smaller and wins. Know the derivation — it is a one-line calculus argument interviewers enjoy.

Product Extension

Batch analytics is offline by nature: “distinct users per session window” over yesterday’s event log, thousands of ad-hoc time ranges — exactly Mo’s shape, and the reorder-for-locality idea is the same one behind sorting scan ranges in columnar engines. Exact distinct-counts feed A/B dashboards where HyperLogLog’s ±2% is unacceptable (billing, fraud thresholds). The sqrt-block warm-up is the mental model for time-series downsampling tiers (raw + per-block rollups) in monitoring systems.

Language/Runtime Follow-ups

  • Python: inline the add/remove logic into the loops (as above) — closures with nonlocal cost ~2× in the hot path; use lists, not defaultdict; at DQUERY limits CPython is marginal, PyPy comfortable.
  • C++: Mo’s home turf — sort with a comparator, int cnt[MAX], easily 10⁷+ pointer moves/sec. Use inline void add(int i).
  • Java: sort an Integer[] of indices or a primitive index array with a manual sort; avoid HashMap for cnt at all costs.
  • Go: sort.Slice on the index slice; flat []int for cnt; closures are cheap enough here.
  • JS/TS: Int32Array for cnt and va; numeric sort comparator (x, y) => keyX - keyY carefully — default sort is lexicographic.

Common Bugs

  1. Shrink before expand: reordering the four while-loops can drive the window through an inverted state (cl > cr+1), sending counts negative and distinct permanently wrong. Expand (cl↓, cr↑) first, always.
  2. Answers in sorted order: writing ans.append(distinct) instead of ans[qi] = distinct — output silently permuted; passes single-query tests, fails everything real.
  3. Block size zero: int(n / √q) truncates to 0 when q > n² — every l lands in “block 0/0” → ZeroDivisionError or a degenerate sort. Clamp with max(1, ·).
  4. Dict instead of compressed array: correct output, 5–10× slowdown in Python — the classic “my algorithm is right but TLEs” report.
  5. Inclusive/exclusive drift: window is inclusive [cl, cr] starting at (0, −1); mixing in half-open query ranges off-by-ones every answer by ±1 on boundary elements.
  6. Sorting by (l, r) without blocks: looks plausible, degrades to O(n·q) pointer travel on adversarial queries — the entire algorithm is the sort order.

Debugging Strategy

First verify the window invariant in isolation: run Mo’s with queries in given order (no sort) on a 10-element array and diff every answer against len(set(a[l:r+1])) — if that fails, the bug is in add/remove or loop order, not the sort. Then re-enable the sort and diff again — a newly appearing mismatch means bug 2 (answer indexing). For performance, count pointer steps with a counter and print it: at n = 3×10⁴, q = 2×10⁵ expect ~10⁷ (≈ (n+q)√n), not ~10⁹; a 100× excess means the sort key is wrong (bug 6) or B is degenerate (bug 3). For SqrtBlocks, test l and r inside the same block — the whole-block loop must not fire.

Mastery Criteria

  • Wrote mo_distinct from cold start in <20 minutes, correct on the DQUERY sample first run.
  • Derived the O((n+q)√n) movement bound on a whiteboard, including the per-block r argument.
  • Stated why distinct-count cannot live in a segment tree (no merge), unprompted.
  • Wrote the sqrt-block struct in <8 minutes.
  • Given a fresh range-query problem, chose correctly between BIT/segment tree, Mo’s, and persistence in <1 minute (online? updates? mergeable?).
  • Named the tuned block size n/√q and the snake/Hilbert orderings as constant-factor upgrades.