Lab 15 — Persistent Segment Tree

Goal

Build a version-preserving segment tree via path copying — each update creates a new O(log n)-node version while sharing the rest — enabling queries against any historical version, and the canonical application: k-th smallest in a range in O(log n) per query.

Background

Path copying: an update to a segment tree touches one root-to-leaf path — O(log n) nodes. Instead of mutating, allocate fresh copies of exactly those nodes; their untouched children are shared with the previous version. Each version is just a root id; all versions stay queryable forever.

The canonical application — k-th smallest in range [l, r]: build version i as a count tree over the compressed value axis after inserting a[0..i−1] (prefix trees). Then cnt(root[r+1]) − cnt(root[l]) at any tree node gives how many elements of a[l..r] fall in that value interval — subtraction of versions works node-by-node because the trees are structurally identical. Walk both roots down simultaneously like an order-statistics tree: if the left subtree difference ≥ k go left, else subtract and go right.

Alternatives (know the trade): Mo’s algorithm (phase-03 lab-12) answers offline in O((n + q)√n) with a frequency structure; a merge sort tree answers via binary search in O(log³ n) per query. Persistence is online, O(log n) per query, and handles “query version t” for free.

Persistence generalizes: persistent DSU (rollback / partial persistence), persistent tries — a persistent binary trie with the same prefix-difference trick solves max-XOR-with-element-in-range.

Interview Context

  • Codeforces Div 1 / ICPC: standard tool; MKTHNUM-style problems are a recognized genre
  • Competitive programming interviews (quant, HFT): plausible
  • Systems design adjacency: MVCC snapshots, functional data structures — concept, not code
  • Standard FAANG: implementation essentially never asked; ordinary segment trees are already rare
  • LC has no exact k-th-smallest-in-range problem; nearest neighbors (LC 315 family) are solved with BIT/merge sort. Anchor is SPOJ MKTHNUM / the Codeforces standard.

When to Skip This Topic

Skip if any of these are true:

  • You aren’t targeting ICPC or competitive-programming-flavored interviews
  • Ordinary segment trees (build/update/query) aren’t automatic yet
  • You haven’t done coordinate compression under time pressure

This is competitive programming infrastructure. For FAANG prep, the transferable 20% is the idea of structural sharing — read Background, skip the drill.

Problem Statement

K-th smallest number in a range (SPOJ MKTHNUM semantics).

Given an array a of n integers and q queries (l, r, k), answer for each: the k-th smallest element among a[l..r] (1-indexed k; k = 1 is the minimum). Queries are online — no batching, no reordering.

Constraints

  • n ≤ 10^5, q ≤ 10^4
  • values up to 10^9 in magnitude, duplicates allowed → coordinate compression mandatory
  • 1 ≤ l ≤ r ≤ n, 1 ≤ k ≤ r − l + 1
  • Memory: O((n + q) log n) nodes — must fit; this kills naive per-version copies

Clarifying Questions

  1. Are queries online or can I batch them? (Online — otherwise Mo’s or offline sweeps compete.)
  2. Duplicates in the array? (Yes — count trees handle them; each occurrence counts separately.)
  3. Is the array immutable after build? (Yes for MKTHNUM; updates change the tool — see follow-ups.)
  4. k is 1-indexed? (Yes; k-th smallest, ties by multiplicity.)
  5. Value range? (Up to 10^9 — compress to ranks, answer maps back through the sorted values.)

Examples

a = [2, 6, 3, 1, 6]           (1-indexed positions 1..5)
query (1, 5, 3): sorted slice [1, 2, 3, 6, 6] → 3
query (2, 4, 2): slice [6, 3, 1] sorted [1, 3, 6] → 3
query (1, 5, 5): → 6   (duplicate 6 counted twice)
query (3, 3, 1): → 3   (single element)

Brute Force

Per query: sorted(a[l-1:r])[k-1]. O(r − l) copy + O(m log m) sort per query. For n = 10^5, q = 10^4 with full-range queries: ~10^9 × log — TLE in any language that isn’t cheating. Quickselect drops the log but keeps the O(n) copy per query.

Brute Force Complexity

  • Time: O(q · n log n)
  • Space: O(n) per query slice

Optimization Path

  1. Compress values to ranks 0..m−1 (sorted uniques); answers map back via the sorted array.
  2. Build version 0 = empty tree over the value axis. Version i = version i−1 with rank(a[i]) inserted (count +1 along one path). Path copying: O(log m) new nodes per version.
  3. For query (l, r, k): the multiset of a[l..r] is version r “minus” version l−1 — valid per node since all versions share shape over the same value axis.
  4. Descend both roots in lockstep: left_count = cnt[left[root_r]] − cnt[left[root_l]]; k ≤ left_count → recurse left, else k −= left_count, recurse right. Leaf index is the answer’s rank.

Final Expected Approach

class PersistentSegTree:
    """Count tree over compressed value axis [0, m). Node 0 = shared empty tree."""

    def __init__(self, m):
        self.m = m
        self.left, self.right, self.cnt = [0], [0], [0]   # node 0: all zeros

    def _node(self, l, r, c):
        self.left.append(l); self.right.append(r); self.cnt.append(c)
        return len(self.cnt) - 1

    def update(self, prev, pos):
        """New version = prev with count at pos incremented. prev is untouched."""
        return self._update(prev, pos, 0, self.m - 1)

    def _update(self, node, pos, lo, hi):
        if lo == hi:
            return self._node(0, 0, self.cnt[node] + 1)
        mid = (lo + hi) // 2
        if pos <= mid:
            return self._node(self._update(self.left[node], pos, lo, mid),
                              self.right[node], self.cnt[node] + 1)
        return self._node(self.left[node],
                          self._update(self.right[node], pos, mid + 1, hi),
                          self.cnt[node] + 1)

    def kth(self, root_l, root_r, k):
        """k-th smallest (1-indexed) in the version difference root_r − root_l."""
        lo, hi = 0, self.m - 1
        while lo < hi:
            mid = (lo + hi) // 2
            in_left = self.cnt[self.left[root_r]] - self.cnt[self.left[root_l]]
            if k <= in_left:
                root_l, root_r = self.left[root_l], self.left[root_r]
                hi = mid
            else:
                k -= in_left
                root_l, root_r = self.right[root_l], self.right[root_r]
                lo = mid + 1
        return lo                                    # rank of the answer

def build_kth_oracle(a):
    """query(l, r, k): k-th smallest of a[l..r], 0-indexed inclusive bounds."""
    vals = sorted(set(a))
    rank = {v: i for i, v in enumerate(vals)}
    pst = PersistentSegTree(len(vals))
    roots = [0]                                      # roots[i] covers a[0..i-1]
    for x in a:
        roots.append(pst.update(roots[-1], rank[x]))

    def query(l, r, k):
        return vals[pst.kth(roots[l], roots[r + 1], k)]

    return query

Note the null-node trick: node 0 is a real all-zero node whose children are itself — descending into unbuilt subtrees needs no None checks, and version 0 (empty tree) is just root 0.

Data Structures

  • Three parallel lists left[], right[], cnt[] — nodes as array slots, not objects
  • roots[]: version i’s root id; the entire “history” is this list of ints
  • Sorted unique values + rank dict for coordinate compression

Correctness Argument

  • Prefix property: version i counts exactly a[0..i−1]; counts are non-negative and monotone in i, so cnt_r − cnt_l at any node = occurrences of that value interval within a[l..r].
  • Shape invariance: every version is a tree over the same interval decomposition of [0, m); node-wise subtraction is therefore meaningful even across shared/unshared nodes.
  • Descent invariant: entering a subtree, k is the rank of the answer within that subtree’s value interval restricted to a[l..r]; the left-count comparison preserves it. At a leaf, interval = single rank = answer.
  • Persistence: _update never writes to an existing slot — old roots are immutable by construction.

Complexity

  • Build: n updates × O(log m) new nodes → O(n log n) time and space
  • Query: O(log m) descent, O(1) extra space
  • Total memory: O((n + q) log n) node slots (q term appears once updates enter the picture)

Implementation Requirements

  • Node 0 must be a genuine empty node (cnt 0, children 0) — no None/sentinel branching
  • Never mutate an existing node; every changed node is a fresh append
  • Coordinate compression built once, from the whole array, before any update
  • Iterative kth (depth ~17 recursion is fine in Python, but iterative is free)
  • Query answers must map ranks back to original values

Tests

  • Worked example above: a = [2, 6, 3, 1, 6], all four queries
  • Randomized stress: 200 random arrays (n ≤ 60, values −50..50) × 50 random (l, r, k) vs sorted(a[l:r+1])[k-1] — exact match
  • Duplicates-heavy stress: values drawn from {0..3}
  • Version isolation: after full build, query old roots (0, roots[i]) for every prefix and every k — history must be intact
  • Scale: n = 10^5 build + 10^4 random queries; measure time and node count (~n log n ≈ 1.7 × 10^6 nodes)
  • Edge: n = 1; l = r; k = 1 and k = r − l + 1

Follow-up Questions

  • Memory analysis? → O(n log n) node slots for the build, O(log n) per subsequent update. In Python, three flat lists of ints beat node objects by a large constant: no per-object header (~56 bytes each), no attribute dicts, better cache behavior; in CPython this is often the difference between AC and MLE.
  • k-th smallest with point updates? → Prefix trees break (one update dirties all later versions). Standard fix: a BIT (Fenwick) indexed by position whose cells are segment-tree roots over the value axis; a query walks O(log n) roots simultaneously → O(log² n) per op.
  • Max XOR with an element in a[l..r]? → Same prefix-difference trick on a persistent binary trie over bits: descend choosing the opposite bit where cnt_r − cnt_l > 0 proves existence in range.
  • Does anything in production actually work like this? → Yes — Clojure and Scala vectors are persistent 32-ary tries (path copying with 32-way branching, ~log₃₂ n depth); git’s object model and MVCC database snapshots are the same structural-sharing idea.
  • Other structures answering k-th-in-range? → Wavelet trees: O(log σ) per query, one structure, no versions. Merge sort tree: O(log³ n), simpler. Fractional cascading speeds repeated binary searches in layered structures. Name-drop level: know they exist and their trade-offs.
  • Offline alternatives? → Sort queries via Mo’s algorithm (phase-03 lab-12) with a value-frequency BIT: O((n + q)√n log n) — worse here, but wins when the per-element update is O(1) and queries are batchable.

Product Extension

  • Time-travel queries: “what was the p95 latency distribution as of version t” over append-only metrics
  • MVCC storage engines: snapshot reads without blocking writers (Postgres, LMDB copy-on-write B-trees)
  • Version control internals: content-addressed trees sharing unchanged subtrees (git)
  • Immutable state in UIs: undo/redo stacks with structural sharing rather than deep copies

Language/Runtime Follow-ups

  • Python: parallel lists as shown; preallocating capacity helps little — append is amortized O(1); PyPy or arrays (array('i')) if TLE.
  • C++: the reference implementation — struct pool with int indices, new is too slow; ~40 bytes/node × 2M nodes is nothing.
  • Java: int arrays sized 4 × (n log n) upfront; avoid Integer boxing in the pool.
  • Clojure/Scala: persistence is the default — but 32-ary tries don’t give you the count-difference descent; you’d still hand-roll this.

Common Bugs

  1. Mutating a shared node: writing cnt[node] += 1 instead of allocating — silently corrupts every earlier version; the version-isolation test exists to catch exactly this.
  2. Prefix off-by-one: the query pair is (roots[l], roots[r + 1]) for 0-indexed inclusive [l, r] — equivalently (roots[l−1], roots[r]) 1-indexed. Mixing conventions gives answers shifted by one element.
  3. Compression per query: ranks must come from the global sorted array; compressing a slice changes the axis and all counts.
  4. Null-node branching: using None children forces checks at every step and breaks the lockstep descent; the self-referential zero node eliminates the whole bug class.
  5. k index base: kth expects 1-indexed k. Feeding a 0-indexed k breaks immediately at the minimum (k = 0 never satisfies k ≤ in_left with in_left = 0 correctly) — always test k = 1 and k = r − l + 1.
  6. Returning the rank, not the value: the descent yields a compressed index; forgetting vals[...] is invisible on arrays that happen to equal their ranks — test with values like 10^9.

Debugging Strategy

  • Query (0, roots[n], k) for k = 1..n and compare against the fully sorted array — validates build + descent without range subtraction
  • Then fix l = 0 and vary r (pure prefixes), then general (l, r) — isolates which of the three layers is broken
  • Assert cnt[left[x]] + cnt[right[x]] == cnt[x] for every internal node after build
  • Print node-pool size after build: far from n·⌈log₂ m⌉ + 1 means paths are wrong or sharing is broken

Mastery Criteria

  • Implement build + update + kth from scratch, bug-free vs stress test, in ≤ 40 min
  • State the memory bound and the node-pool-vs-objects argument in ≤ 3 min
  • Explain why version subtraction is valid (shape invariance) without prompting
  • Adapt the prefix-difference trick to a persistent trie (max-XOR in range) in ≤ 30 min
  • Choose correctly between persistence, Mo’s, merge sort tree, and wavelet tree given constraints in ≤ 5 min