Lab 01 — Segment Tree (Point Update + Range Query)

Goal

Implement a segment tree from scratch that supports point updates and range-sum / range-min / range-max queries on an array of N integers. Build in O(N), query and update in O(log N) each. Internalize the recursion structure so you can re-derive any aggregate variant on the fly. After this lab you should be able to write a working segment tree from blank slate in under 12 minutes with zero off-by-ones.

Background Concepts

A segment tree represents an array as a near-balanced binary tree where each internal node stores the aggregate (sum / min / max / gcd / …) of a contiguous range. Leaves correspond to individual array elements; internal nodes correspond to the union of their children’s ranges. The tree has depth O(log N), so any range [l, r] decomposes into at most 2 · log₂(N) disjoint subtree-ranges. That is the entire complexity argument: each query and update walks O(log N) nodes.

The tree is conventionally stored in a flat array of size 4N (worst-case nearly-balanced binary tree on N leaves) with the root at index 1, left child at 2 · i, right child at 2 · i + 1. This avoids pointer overhead and is cache-friendly.

The aggregate must be associative so that subtree results can be combined. Sum, min, max, gcd, xor, “and”, “or”, and matrix multiplication all qualify. Median, mode, and “k-th smallest” do not combine cleanly and need different structures.

Interview Context

Range queries with updates appear in 3–5% of FAANG-tier Hard pools, but they appear more often on the bar-raiser round. Recognizing that prefix sums (O(1) query, O(N) update) won’t survive the workload — that you need O(log N) for both — is the reflex this lab builds. Companies that screen with segment trees: Meta (frequent), Google (occasional), Amazon (rare), Stripe / HFT shops (very frequent — order book, sliding aggregates). Bombing this is a no-hire signal at L5+.

Problem Statement

Implement a class NumArray initialized with an integer array nums. Support:

  • update(i, val): set nums[i] = val.
  • sumRange(left, right): return the sum of nums[left..right] inclusive.

Both must be O(log N). After implementing the sum variant, refactor so swapping combine = + for combine = min or combine = max requires changing one line.

Constraints

  • 1 ≤ N ≤ 3 × 10⁴
  • −100 ≤ nums[i] ≤ 100
  • 0 ≤ i, left, right < N
  • Up to 3 × 10⁴ calls to update and sumRange combined.

Clarifying Questions

  1. Are queries inclusive on both ends? (Yes — [left, right].)
  2. Is nums mutable in place, or owned by NumArray? (Owned; copy on construction.)
  3. Are the values guaranteed to fit in int32? (Yes; sum across N elements at value ±100 fits comfortably.)
  4. Is update an assignment or delta? (Assignment — set, not add.)

Examples

NumArray([1, 3, 5])
sumRange(0, 2) → 9
update(1, 2)        // array becomes [1, 2, 5]
sumRange(0, 2) → 8
update(0, 10)       // array becomes [10, 2, 5]
sumRange(1, 2) → 7

Initial Brute Force

Store nums as a plain list. update(i, v): nums[i] = v (O(1)). sumRange(l, r): sum(nums[l:r+1]) (O(N)). Updates are fast; queries are linear.

Brute Force Complexity

Update O(1). Query O(N). Total over Q queries + U updates: O(N · Q + U). At N = Q = 3 × 10⁴: 9 × 10⁸ ops — TLE.

Optimization Path

Two natural alternatives.

Prefix sums. prefix[i] = nums[0] + ... + nums[i-1]. Query is prefix[r+1] - prefix[l] in O(1). But update(i, v) requires recomputing prefix[i+1..N] in O(N). Wrong tradeoff for this workload.

Sqrt decomposition. Block size √N; per-block sums. Update O(1), query O(√N). Total O(N · √N) = O(N^1.5) — at N = 3 × 10⁴, ~5 × 10⁶ ops. Passes but is suboptimal and crusty.

Segment tree. Build O(N), update O(log N), query O(log N). Total O((N + Q) log N) = ~5 × 10⁵ ops. Clean fit.

Final Expected Approach

Recursive segment tree on a flat array of size 4N.

  • Build build(node, nl, nr): if nl == nr, leaf = arr[nl]; else recurse left and right, set tree[node] = tree[2node] + tree[2node + 1].
  • Update update(node, nl, nr, idx, val): recurse into the child whose range contains idx; on return, recompute parent.
  • Query query(node, nl, nr, ql, qr): total miss → identity (0 for sum, +∞ for min); total cover → return tree[node]; partial → recurse both children and combine.

Public API wraps with node = 1, nl = 0, nr = N - 1.

Data Structures Used

  • A single integer array tree[] of size 4N (sum aggregates).
  • Optional integer n storing the original length.

Correctness Argument

Build establishes the invariant tree[node] = combine over [nl, nr] by induction on subtree size. Update: along the recursion, only nodes whose range contains idx are touched; each is recomputed from its (now-correct) children, so the invariant is preserved. Query decomposes [ql, qr] into a disjoint union of subtree ranges; the result is the combine of those. The decomposition has size ≤ 2 log N because along any root-to-leaf path the recursion either stops (full cover or miss) or splits at most twice (once for the left boundary, once for the right). The total work is O(log N).

Complexity

OperationTimeSpace
BuildO(N)O(4N)
UpdateO(log N)O(log N) recursion
QueryO(log N)O(log N) recursion

Implementation Requirements

class NumArray:
    def __init__(self, nums):
        self.n = len(nums)
        self.tree = [0] * (4 * self.n)
        self._build(1, 0, self.n - 1, nums)

    def _build(self, node, nl, nr, a):
        if nl == nr:
            self.tree[node] = a[nl]; return
        mid = (nl + nr) // 2
        self._build(2*node, nl, mid, a)
        self._build(2*node + 1, mid + 1, nr, a)
        self.tree[node] = self.tree[2*node] + self.tree[2*node + 1]

    def update(self, i, val):
        self._update(1, 0, self.n - 1, i, val)

    def _update(self, node, nl, nr, idx, val):
        if nl == nr:
            self.tree[node] = val; return
        mid = (nl + nr) // 2
        if idx <= mid: self._update(2*node, nl, mid, idx, val)
        else: self._update(2*node + 1, mid + 1, nr, idx, val)
        self.tree[node] = self.tree[2*node] + self.tree[2*node + 1]

    def sumRange(self, l, r):
        return self._query(1, 0, self.n - 1, l, r)

    def _query(self, node, nl, nr, ql, qr):
        if qr < nl or ql > nr: return 0           # miss
        if ql <= nl and nr <= qr: return self.tree[node]   # cover
        mid = (nl + nr) // 2
        return self._query(2*node, nl, mid, ql, qr) + self._query(2*node + 1, mid + 1, nr, ql, qr)

Refactor to support min by changing the identity (+∞), the leaf assignment (still a[nl]), and the combine (min).

Tests

  • N=1: update(0, 5); sumRange(0, 0) == 5.
  • All zeros: every range query returns 0.
  • All same value: sumRange(l, r) == val * (r - l + 1).
  • After update(i, v): sumRange(i, i) == v; sum across full range matches direct sum of array.
  • Random fuzz: 1000 ops alternating updates and queries against a brute-force list.
  • Min variant: build [3, 1, 4, 1, 5, 9, 2, 6], query(2, 5) == 1; after update(3, 10), query(2, 5) == 4.

Follow-up Questions

  1. “Now I want range updates.” → Lab 02 (lazy propagation).
  2. “Now I want O(1) queries on a static array.” → sparse table (Lab 04).
  3. “Now the array is 2D.” → segment tree of segment trees, O(log² N) per op.
  4. “Make it iterative.” → power-of-two-padded leaves at indices [N, 2N); update walks i // 2 upward, query walks l, r toward the middle.
  5. “How would you support count of values ≥ K in [l, r]?” → merge sort tree (segment tree with sorted lists) or wavelet tree.

Product Extension

Real-time analytics dashboards: a stream of N metrics with both edits and arbitrary-range aggregates (e.g., “total revenue from days 17–24” while a correction is being applied to day 19). The naive list is fine until the workload has both fast updates and fast arbitrary-range queries on the same data — then a segment tree over the time axis, keyed by index, is what powers the underlying store.

Language/Runtime Follow-ups

  • Python: recursion at N=3×10⁴ is fine but the per-op constant is high. For larger N convert to iterative or use sys.setrecursionlimit. Consider array.array('i', ...) over plain list for cache locality.
  • Java: use int[] tree = new int[4 * n]; the 4n allocation is critical because computing next_power_of_two(n) * 2 is a fencepost-error magnet. Method dispatch has a real cost — inline the recursion if hot.
  • Go: no recursion-limit issue; keep the slice. Use int (sized to platform) unless the problem demands int64.
  • C++: the canonical implementation. vector<long long> tree(4 * n). Inline the body; mark methods inline. For competitive problems use the iterative version (template by Adrian Panaete or Codeforces “EDU”).
  • JS/TS: typed arrays — new Int32Array(4 * n) — outperform plain arrays. Recursion depth at N=3×10⁴ is fine in V8.

Common Bugs

  1. Mixing closed-interval [ql, qr] with half-open [ql, qr) between query and recursion. Always pick closed and stay consistent.
  2. Sizing tree as 2N instead of 4N: works for power-of-two N, segfaults otherwise.
  3. Forgetting to recompute tree[node] after recursing in update. The leaf updates correctly but the parent stays stale.
  4. Identity wrong for the aggregate: 0 for sum, but +∞ (float('inf') / Long.MAX_VALUE) for min and −∞ for max. Returning 0 from a min query missing-range gives wrong answers silently.
  5. Building from nums[mid + 1] vs nums[mid] — pick one slicing convention.
  6. Iterative version: forgetting to round N up to a power of two before placing leaves.

Debugging Strategy

When tests fail, drop into a tiny instance (N=4, indices 0..3) and print(tree) after each op. Verify by hand: tree[1] should equal sum over [0,3], tree[2] over [0,1], tree[3] over [2,3]. If those don’t hold post-build, your build recursion is broken — fix that before touching update or query. Add assert for the cover/miss/partial branches printing (node, nl, nr, ql, qr) to spot which sub-call returns the wrong total.

Mastery Criteria

  • Recognized the segment-tree signal in <60 seconds from a “range query + point update” problem statement.
  • Wrote build/update/query on a blank screen in <12 minutes with no off-by-ones, first try.
  • Refactored sum → min → max in <2 minutes by changing identity + combine only.
  • Stated complexity (build O(N), update/query O(log N), space O(4N)) without prompting.
  • Solved LC 307 in <15 minutes from cold start.
  • Solved one cousin problem (LC 308 or LC 1157) in <30 minutes from cold start.