cswork / libraryOpen practice workspace

Algorithm tutorials

Segment Tree

A binary tree of range summaries that answers range queries and point updates in O(log n) for any mergeable operation. Includes a compact bottom-up template, lazy propagation for range updates, and tree descent for 'first index' queries.

Intuition

A segment tree is a binary tree built over an array. Every node is responsible for a contiguous range and stores a summary of it: the sum, the minimum, the maximum, a gcd, or a small custom record. The root covers the whole array, each internal node splits its range in half between its two children, and the leaves are single elements. The height is about log2(n).

Two facts make it work:

  • Any range is a union of O(log n) nodes. To answer a query on [l, r), take the maximal nodes that fit entirely inside it; there are at most two per level. Combine their summaries from left to right.
  • A point update touches one root-to-leaf path. Change the leaf, then recompute each ancestor from its two children. That is O(log n) nodes.

The summary must be mergeable: an associative operation with an identity element (sum with 0, min with +infinity, max with -infinity, gcd with 0). The operation does not need to be commutative, as long as you always merge the left part before the right part. That is what makes richer summaries possible: for example, storing (total, best prefix, best suffix, best subarray) per node answers "maximum subarray sum inside [l, r)" with updates.

Compared with a Fenwick tree, a segment tree handles non-invertible operations (min, max, gcd), range updates other than adding to sums (range assign, range add with range max), and tree descents on any summary, at the cost of more code. If you only need sums with point updates, the Fenwick tree is shorter; if the array never changes, prefix sums or a sparse table are simpler.

When to use it in interviews / OAs

  • Range min/max/gcd with point updates: a mutable range-minimum query, the best price in a time window with corrections.
  • Range updates plus range queries: "add v to every element in [l, r)", then sums or maxima. This needs lazy propagation (below).
  • Searching by a condition: the first index whose value is at least x, the first position where a prefix sum reaches k, "first seat block with enough free seats". Walking down the tree answers in O(log n) instead of binary search plus queries in O(log^2 n).
  • Interval scheduling on coordinates: maximum number of overlapping bookings, falling squares (range assign, range max), rectangle area union (sweep line with coverage counts). Compress coordinates first.
  • Values as positions: a tree over compressed values works as an order-statistics multiset (count less than x, k-th smallest), like a Fenwick tree over values.

Template

A compact bottom-up (iterative) segment tree: leaves at indices n .. 2n-1, node i has children 2i and 2i + 1. It works for any n, not only powers of two, and keeps separate left and right accumulators so even non-commutative operations combine in order.

from typing import Callable, Generic, TypeVar

T = TypeVar("T")


class SegmentTree(Generic[T]):
    # Point assignment + range query for any associative op with an identity.
    def __init__(self, values: list[T], op: Callable[[T, T], T], identity: T) -> None:
        self.n, self.op, self.e = len(values), op, identity
        self.t = [identity] * self.n + list(values)  # Leaves live at n .. 2n-1.
        for i in range(self.n - 1, 0, -1):  # Fill parents bottom-up; node i has children 2i, 2i+1.
            self.t[i] = op(self.t[2 * i], self.t[2 * i + 1])

    def update(self, i: int, value: T) -> None:
        i += self.n
        self.t[i] = value
        while i > 1:
            i //= 2
            self.t[i] = self.op(self.t[2 * i], self.t[2 * i + 1])

    def query(self, l: int, r: int) -> T:
        # op over values[l:r], combined strictly left to right.
        left, right = self.e, self.e
        l += self.n
        r += self.n
        while l < r:
            if l & 1:  # l is a right child: take it, then step past it.
                left = self.op(left, self.t[l])
                l += 1
            if r & 1:  # r is exclusive, so the node r - 1 is fully inside.
                r -= 1
                right = self.op(self.t[r], right)
            l //= 2
            r //= 2
        return self.op(left, right)


st = SegmentTree([5, 2, 8, 1, 9], min, float("inf"))
assert st.query(0, 3) == 2 and st.query(2, 5) == 1
st.update(3, 7)
assert st.query(2, 5) == 7 and st.query(1, 2) == 2

Passing op and identity makes it reusable: (min, inf), (max, -inf), (operator.add, 0), (math.gcd, 0), or your own merge function for tuples.

Extension: range updates with lazy propagation

Updating every element in a range one by one costs O(n log n). Lazy propagation applies the update only to the O(log n) nodes that exactly cover the range: update their summaries immediately and leave a tag that says "my children still owe this update". Before a later operation descends below a tagged node, it pushes the tag one level down. Every operation still touches O(log n) nodes.

The recursive version below supports range add, range sum and range max, plus a descent that finds the leftmost value >= x: at each node, go left if the left child's maximum is large enough, otherwise go right.

class LazySegmentTree:
    # Range add, range sum, range max, and "first index with value >= x".
    def __init__(self, values: list[int]) -> None:
        self.n = len(values)
        size = 4 * max(1, self.n)  # Recursive layout needs up to 4n nodes.
        self.sm = [0] * size
        self.mx = [0] * size
        self.tag = [0] * size  # Pending addition owed to the children.
        if self.n:  # An empty array has no root to build.
            self._build(1, 0, self.n - 1, values)

    def _build(self, node: int, lo: int, hi: int, values: list[int]) -> None:
        if lo == hi:
            self.sm[node] = self.mx[node] = values[lo]
            return
        mid = (lo + hi) // 2
        self._build(2 * node, lo, mid, values)
        self._build(2 * node + 1, mid + 1, hi, values)
        self._pull(node)

    def _pull(self, node: int) -> None:
        self.sm[node] = self.sm[2 * node] + self.sm[2 * node + 1]
        self.mx[node] = max(self.mx[2 * node], self.mx[2 * node + 1])

    def _apply(self, node: int, lo: int, hi: int, v: int) -> None:
        self.sm[node] += v * (hi - lo + 1)  # The sum grows by v times the length.
        self.mx[node] += v
        self.tag[node] += v

    def _push(self, node: int, lo: int, hi: int) -> None:
        if self.tag[node]:  # Hand the pending addition down before visiting children.
            mid = (lo + hi) // 2
            self._apply(2 * node, lo, mid, self.tag[node])
            self._apply(2 * node + 1, mid + 1, hi, self.tag[node])
            self.tag[node] = 0

    def add(self, l: int, r: int, v: int) -> None:  # values[l:r] += v
        def go(node: int, lo: int, hi: int) -> None:
            if r <= lo or hi < l:
                return
            if l <= lo and hi < r:  # Fully covered: update this node, stop descending.
                self._apply(node, lo, hi, v)
                return
            self._push(node, lo, hi)
            mid = (lo + hi) // 2
            go(2 * node, lo, mid)
            go(2 * node + 1, mid + 1, hi)
            self._pull(node)

        go(1, 0, self.n - 1)

    def query(self, l: int, r: int) -> tuple[int, int]:  # (sum, max) of values[l:r], l < r
        def go(node: int, lo: int, hi: int) -> tuple[int, int]:
            if r <= lo or hi < l:
                return 0, float("-inf")
            if l <= lo and hi < r:
                return self.sm[node], self.mx[node]
            self._push(node, lo, hi)
            mid = (lo + hi) // 2
            s1, m1 = go(2 * node, lo, mid)
            s2, m2 = go(2 * node + 1, mid + 1, hi)
            return s1 + s2, max(m1, m2)

        return go(1, 0, self.n - 1)

    def first_at_least(self, x: int) -> int:  # Leftmost index with value >= x, or -1.
        def go(node: int, lo: int, hi: int) -> int:
            if self.mx[node] < x:
                return -1  # Nothing in this subtree qualifies.
            if lo == hi:
                return lo
            self._push(node, lo, hi)
            mid = (lo + hi) // 2
            res = go(2 * node, lo, mid)  # Prefer the left half.
            return res if res != -1 else go(2 * node + 1, mid + 1, hi)

        return go(1, 0, self.n - 1) if self.n else -1


lt = LazySegmentTree([1, 3, 2, 0, 4])
lt.add(1, 4, 2)  # values = [1, 5, 4, 2, 4]
assert lt.query(0, 5) == (16, 5) and lt.query(2, 4) == (6, 4)
assert lt.first_at_least(4) == 1 and lt.first_at_least(6) == -1

What changes for other update types: the tag must say how to update a node's summary without visiting its children (adding v to a range changes its sum by v * length), and two tags must be composable. For "assign v to the range", a new assignment overwrites any older tag, and you need a sentinel such as None for "no pending assignment", because 0 is a valid value. Combining add and assign is possible but needs a careful order of operations.

Huge index ranges. If positions go up to 10^9, compress the coordinates that actually occur (offline), or create nodes only when they are first touched (a dynamic segment tree, with children stored in dicts or growing arrays).

A note on persistence. A persistent structure keeps every old version queryable. In interviews this usually means snapshots of an array or map, solved by storing a per-key history of (version, value) pairs and binary searching it (snapshot arrays, time-based key-value stores). Segment trees can be made persistent by path copying: an update copies only the O(log n) nodes on its root-to-leaf path and shares everything else with the previous version. That powers contest techniques such as k-th smallest in an arbitrary subarray, and is beyond interview scope.

Complexity

OperationTime
buildO(n)
point updateO(log n)
range queryO(log n)
range update with lazy tagsO(log n)
descent ("first index with ...")O(log n)

Space: 2n entries for the bottom-up layout, up to 4n per array for the recursive layout.

Constant factors matter in Python. On a typical laptop, the bottom-up tree above runs 10^5 updates plus 10^5 queries on 10^5 elements in under a second, while the recursive lazy tree needs roughly 3 seconds per 10^5 operations. For heavy OA inputs, prefer the iterative tree, or a Fenwick tree when it is enough.

Common pitfalls

  • Array too small. The recursive layout needs 4n slots; 2n is only enough for the bottom-up layout.
  • Wrong identity. Using 0 for a min tree silently returns 0 for ranges of positive numbers. Use inf for min, -inf for max.
  • Merge order. For non-commutative summaries, always combine left then right, including in the query accumulators.
  • Forgetting push or pull. Push tags before visiting children (in updates, queries and descents) and recompute the parent after updating children.
  • Lazy tag semantics. Decide whether a tagged node's own summary is already updated (here it is) and stick to it. For range add on sums, multiply by the segment length.
  • Inclusive vs exclusive bounds. Pick one convention (half-open [l, r) here) and convert once at the boundary.
  • Overkill. Static data with range-sum queries only needs prefix sums; mention the simpler option first in an interview.

Interview variations

  • Range sum query, mutable and range minimum query with updates.
  • My calendar I/II/III: count overlaps with range add and range max over compressed or dynamic coordinates.
  • Falling squares: range assign and range max.
  • Booking concert tickets in groups: sums for capacity checks plus a descent for the first row with enough free seats.
  • Longest increasing subsequence with a bounded difference: a max tree over values, point updates as you scan.
  • Maximum subarray sum on a range with updates: four values per node.
  • Count of smaller numbers after self: a counting tree over compressed values (a Fenwick tree is shorter).