Lab 02 — Segment Tree With Lazy Propagation

Goal

Extend Lab 01’s segment tree to support range updates in O(log N) using lazy propagation. Implement range-add + range-sum and articulate the push-down invariant so cleanly that you can re-derive lazy on a different aggregate (range-set + range-sum, range-flip + range-count) under interview pressure.

Background Concepts

Lazy propagation defers work. When an update covers a whole subtree, instead of recursing into all O(2^depth) descendants, you stamp a single lazy tag on that subtree’s root and update its aggregate in O(1). The descendants stay stale until something forces a deeper visit; at that point you push_down the tag — apply it to the children’s aggregates and merge into their lazy slots — and clear the parent’s tag.

This works whenever:

  1. The update operation has an O(1) batch form: applying “add v to all of [nl, nr]” to tree[node] is tree[node] += v * (nr - nl + 1).
  2. The lazy tags compose: a pending “add 3” followed by a new “add 5” composes to “add 8”. Without composition, you cannot stack tags; you must push first.
  3. There is a identity lazy value (e.g., 0 for add) meaning “nothing pending”.

For mixed update types (“add” and “set” both), composition needs an explicit rule: a new “set” wipes any pending “add”; a new “add” composes with a pending “set” by changing the set value.

Interview Context

Asked at: companies with high-frequency-trading or analytics flavor (Stripe, Two Sigma, Jane Street), and Meta in bar-raiser slots. Most interview problems that need this dress up as “support add v to a range and report sum of a range” or as a count-of-overlapping-intervals problem like LC 732. Failing to know this structure caps you at Mediums; recognizing it and implementing it correctly is a green-light at L5+.

Problem Statement

Implement a class RangeArray over n integers (initially zero) supporting:

  • add(l, r, v): add v to every index in [l, r].
  • sumRange(l, r): return the sum of arr[l..r].

Both O(log n).

Constraints

  • 1 ≤ n ≤ 10⁵
  • 1 ≤ Q ≤ 10⁵ ops total.
  • −10⁴ ≤ v ≤ 10⁴.
  • Sums fit in 64-bit (max |sum| ≈ 10⁵ · 10⁵ · 10⁴ = 10¹⁴).

Clarifying Questions

  1. Endpoints inclusive? (Yes.)
  2. Is add cumulative or assignment? (Cumulative — additive.)
  3. Should arr be mutable in place at the leaves? (Conceptually yes; in practice the segment tree owns it.)
  4. 0-indexed or 1-indexed externally? (0-indexed.)

Examples

RangeArray(5)
add(0, 2, 3)         // [3, 3, 3, 0, 0]
sumRange(0, 4) → 9
add(1, 3, 2)         // [3, 5, 5, 2, 0]
sumRange(2, 4) → 7
sumRange(0, 0) → 3

Initial Brute Force

Plain list. add(l, r, v) → for i in range(l, r+1): arr[i] += v (O(N)). sumRange → sum(arr[l:r+1]) (O(N)). Combined per-op O(N).

Brute Force Complexity

O(N) per op. Total O(N · Q) = 10¹⁰ at the limits. TLE by 4 orders of magnitude.

Optimization Path

Difference array? diff[l] += v; diff[r+1] -= v is O(1) per add, but you can only query the final prefix sum after all updates — not interleaved with sum queries. Doesn’t survive mixed workload.

Two Fenwick trees (BIT-RU + BIT-PQ)? Yes, this works for range-add + range-sum specifically — the BIT² trick. Slightly faster constants than segment tree, but only handles invertible aggregates. Segment tree generalizes to range-set, range-min-after-add, range-affine, etc.

Lazy segment tree is the canonical answer.

Final Expected Approach

Augment Lab 01’s tree with a parallel lazy[] array of size 4N, all initialized to 0 (the identity for add).

  • push_down(node, nl, nr): if lazy[node] != 0, apply it to both children’s aggregates (tree[child] += lazy[node] * child_len) and compose it into lazy[child] += lazy[node]. Then clear lazy[node] = 0. Called at the start of any non-leaf update or query that recurses into children.
  • update(node, nl, nr, ql, qr, v):
    • If qr < nl or ql > nr: return (miss).
    • If ql <= nl and nr <= qr: stamp tree[node] += v * (nr - nl + 1); lazy[node] += v; return.
    • Else: push_down; recurse both children; tree[node] = tree[left] + tree[right].
  • query(node, nl, nr, ql, qr): identical structure, with push_down before recursing.

Data Structures Used

  • tree[] — sum aggregates, size 4N, int64.
  • lazy[] — pending add tags, size 4N, int64.

Correctness Argument

Invariant: for every node, tree[node] equals the correct aggregate over its range as if all pending lazy tags up to and including this node have been applied. Specifically, tree[node] is correct; tree[child] may be stale by exactly lazy[node] * child_len.

push_down repairs the children: it adds the missing contribution to their aggregates and composes the tag into theirs (so their own descendants will, later, be repaired similarly). It then clears lazy[node]. After push_down, tree[node] is unchanged and the children are now correct, so descendants of children may be stale only by the children’s own pending tags.

update either (a) misses, doing nothing, (b) totally covers, applying the O(1) batch update directly to tree[node] and stamping the tag, or (c) partially overlaps, which requires push_down before recursing so the children are correct, then recomputes tree[node] from now-current children.

query symmetric.

Complexity

OperationTimeSpace
BuildO(N)O(4N) tree + O(4N) lazy
Range updateO(log N)O(log N) recursion
Range queryO(log N)O(log N) recursion

Implementation Requirements

class RangeArray:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (4 * n)
        self.lazy = [0] * (4 * n)

    def _push_down(self, node, nl, nr):
        if self.lazy[node]:
            mid = (nl + nr) // 2
            left, right = 2*node, 2*node + 1
            self.tree[left] += self.lazy[node] * (mid - nl + 1)
            self.lazy[left] += self.lazy[node]
            self.tree[right] += self.lazy[node] * (nr - mid)
            self.lazy[right] += self.lazy[node]
            self.lazy[node] = 0

    def add(self, l, r, v):
        self._add(1, 0, self.n - 1, l, r, v)

    def _add(self, node, nl, nr, ql, qr, v):
        if qr < nl or ql > nr: return
        if ql <= nl and nr <= qr:
            self.tree[node] += v * (nr - nl + 1)
            self.lazy[node] += v
            return
        self._push_down(node, nl, nr)
        mid = (nl + nr) // 2
        self._add(2*node, nl, mid, ql, qr, v)
        self._add(2*node + 1, mid + 1, nr, ql, qr, v)
        self.tree[node] = self.tree[2*node] + self.tree[2*node + 1]

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

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

Tests

  • N=1, single index: add(0, 0, 5); sumRange(0, 0) == 5.
  • All-zero: any sumRange before any add returns 0.
  • Disjoint adds: add(0, 2, 1), add(3, 5, 2); sumRange(0, 5) == 3 + 6 = 9.
  • Overlapping adds: add(0, 4, 1), add(2, 4, 1); sumRange(2, 4) == 2 + 2 + 2 = 6.
  • Stress: 10⁴ random adds + queries against a brute-force list.
  • Stack of pending tags: add(0, n-1, 1) 100 times; sumRange(i, i) == 100 for all i.

Follow-up Questions

  1. “Now add becomes set (assignment, not delta).” → identity = sentinel None; on push_down, if lazy_parent != None, replace child’s tree and replace child’s lazy.
  2. “Both add and set operations.” → two lazy slots. Composition: a new “set” wipes pending “add”; a new “add” applied while a “set” is pending modifies the set value.
  3. “Range-flip on a binary array, with range-count-of-ones.” → tree[node] = count of 1s; flip → tree[node] = (nr-nl+1) - tree[node]; lazy is a boolean toggle.
  4. “Range-affine: replace a[i] with b · a[i] + c.” → lazy holds (b, c); composition: (b₂, c₂) ∘ (b₁, c₁) = (b₂ b₁, b₂ c₁ + c₂).

Product Extension

A live spreadsheet with array formulas — =ARRAYFORMULA(A1:A1000 + 5) — is exactly range-add. With range-set you get fill-down. With range-affine you get scaling formulas. The backend has to support thousands of these per second per spreadsheet; lazy segment trees are one viable engine.

Language/Runtime Follow-ups

  • Python: 4 × 10⁵ allocations are slow; warm the lists once, never resize. Recursion at N=10⁵ depth ≈ 17, fine.
  • Java: long[] tree, lazy to avoid sum overflow at the limits. Synchronization-free for single-thread; LongAdder is unrelated.
  • Go: same template; []int64 slices.
  • C++: the canonical use case. vector<long long>. Compile with -O2; benchmark on N=10⁶.
  • JS/TS: BigInt64Array is heavy; if values fit in Number’s 53-bit safe range, use Float64Array despite being float (the IEEE 754 representation is exact for ±2⁵³ integers).

Common Bugs

  1. Forgetting push_down before recursing on partial cover. The leaf updates correctly but its sibling subtree returns stale aggregates on later queries. Manifests as queries that depend on update order.
  2. Push down on full-cover branch — wasted work but not wrong; only push on partial overlap.
  3. Identity confusion: 0 is identity for add but a legal value for set. Use None or a sentinel for set-style lazy.
  4. Composition direction: when stamping a new tag onto a parent that already has a tag, write the rule down before coding. For add it’s commutative; for set it isn’t.
  5. int overflow in Java/C++ — sums of up to 10⁵ · 10⁴ values ≈ 10⁹, doubles to 10¹⁴ with adds. Use 64-bit.
  6. Calling push_down on a leaf — guard if nl != nr.

Debugging Strategy

Add an assert_consistent() helper that walks the tree and verifies, for every internal node, tree[node] == tree[left] + tree[right] + lazy[node] * (nr - nl + 1). Wait — that’s not quite the invariant, since lazy[node] has not been pushed yet but tree[node] already includes it. The correct invariant is tree[node] == tree[left] + tree[right] + (lazy[node] * (nr - nl + 1)) only if you treat children’s tree as “before this lazy stamp”. An easier debug helper: after each op, force a full push_down from root to leaves and rebuild aggregates; compare against the brute-force array. If they diverge, you have a push-down-order bug.

Mastery Criteria

  • Stated the push-down invariant in one sentence.
  • Wrote range-add + range-sum lazy seg tree from scratch in <20 minutes, first try.
  • Adapted the same template to range-set + range-sum in <10 additional minutes.
  • Solved LC 732 (My Calendar III) using a coord-compressed lazy seg tree.
  • Stated when not to use lazy (single-point updates → no benefit; non-composing operations → impossible).
  • Pinpointed the canonical bug (missing push_down) within 5 minutes of seeing a failing test.