Segment tree with lazy propagation

Why this matters when you build with AI

Ask an AI tool for a segment tree that supports range updates and you will often get code that looks complete. It has a push function, two recursive calls and a sum at each node. The bugs are in the order of operations. The tool applies a new update to a node, then forgets to push the old tag down, so a later query reads stale children. Or it composes two updates in the wrong order, which gives the right answer for one update and the wrong answer after two.

These bugs do not crash. They give a number that is off, on exactly the operation sequence the sample did not contain. The check that finds them is cheap. Keep a plain array next to the tree, apply every update to both, compare every query, and run a few thousand small random sequences. The rest of this lesson builds the tree with that check in place.

Every program below was compiled with g++ -O2 -std=c++17 -Wall -Wextra on the machine that wrote this lesson, which uses g++ 15.2.0. The random tests use a fixed seed, so you can repeat them. Nothing here was compared with an outside source. The claims are checked by the programs.

The problem this solves

A segment tree answers a range query in O(log n) time. Our tree stores one node per segment of the array. A node holds the sum of its segment. A point update walks one path from the root to a leaf and fixes the sums on that path.

A range update is different. Suppose you must add 5 to every value in [l, r). A plain segment tree would touch every node in the range, which is O(n) in the worst case. Lazy propagation avoids that. When a node's whole segment is covered by the update, the node is updated right away and the change is recorded as a tag. The children are not touched. They learn about the change later, only if a query or an update reaches them.

The word "lazy" means exactly this: work is postponed until something needs the answer. A range update then costs O(log n), the same as a point update.

The idea, with a worked example

Use half-open ranges, the same convention as the prefix sums lesson. The range [l, r) holds the indices l up to r - 1.

Take the array [1, 2, 3, 4, 5] and the range [1, 4), which holds the values 2, 3 and 4. Add 10 to each one. The result is [1, 12, 13, 14, 5].

In the tree, take a node that covers [1, 4) (an example node; a real tree over 5 cells splits at other points). It has sum 2 + 3 + 4 = 9 before the update. After the update its sum is 9 + 3 * 10 = 39. The node stores 39 directly, and it stores a tag "add 10" that says its children still need +10 each. The children of [1, 4) still hold 2, 3 and 4 until someone pushes the tag down.

The rule that makes this correct is the invariant:

A node's stored sum is correct for its own segment, including every tag that was applied to the node itself.
A node's tag is a change that has been applied to the node's sum but not yet to its two children.
To read or change anything inside a node, push its tag to the children first. Pushing a tag means applying it to each child, then clearing it at the node.

If you follow these three rules, every query reads correct numbers.

Affine tags

The tag can be more than "add". Represent every update as an affine map, a function of the form x -> a * x + b. Here a is a multiplier and b is an offset.

Add v is a = 1, b = v.
Assign c is a = 0, b = c, because 0 * x + c is always c.
Multiply by m is a = m, b = 0.

For a node whose segment has len elements and sum s, applying the map (a, b) to every element changes the sum like this:

new sum = a * s + b * len

Check it with the example. The node [1, 4) has s = 9 and len = 3. The map is add 10, so a = 1 and b = 10. The new sum is 1 * 9 + 10 * 3 = 39. That matches the value above.

Composing tags, and why the order matters

A node may receive two maps before it pushes. Suppose the tag is f and a new map g arrives. The combined tag must do what f does and then what g does. Let f = (f.a, f.b) and g = (g.a, g.b). Then:

g(f(x)) = g.a * (f.a * x + f.b) + g.b = (g.a * f.a) * x + (g.a * f.b + g.b)

So the composed tag is (g.a * f.a, g.a * f.b + g.b). Note the order. The old tag f comes first and the new one g comes second.

Here is why the order matters. Start with [1, 2, 3, 4, 5]. Apply two updates to the whole array, first "add 10" and then "assign 7". Each value first becomes its value plus 10, and then every value becomes 7. Now reverse the order. "Assign 7" first makes every value 7, then "add 10" makes every value 17. The two results differ, so a tag that composes in the wrong order gives the wrong sum.

The order check is in the worked problem below, and it prints both orders on the real array.

The standard technique in C++

The code has two parts. The first is a header with the tree. The second reads a problem from standard input. The tree uses 1-based node numbers, with the children of node k at 2k and 2k + 1. It allocates 4n nodes, which is enough for any n.

Values are kept modulo MOD = 998244353. A prime modulus keeps numbers bounded, so a product of two values fits in a signed 64-bit integer.

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
 
const ll MOD = 998244353;
 
// An affine map x -> a*x + b (mod MOD).
struct Affine {
    ll a, b;
};
 
const Affine IDENTITY = {1, 0};
 
// Compose two maps: first apply f, then apply g.
// g(f(x)) = g.a * (f.a * x + f.b) + g.b = (g.a * f.a) * x + (g.a * f.b + g.b)
Affine then(Affine f, Affine g) {
    return {g.a * f.a % MOD, (g.a * f.b + g.b) % MOD};
}
 
// Segment tree over [0, n) with lazy affine tags and range sum.
// Ranges are half-open: [l, r).
struct LazySeg {
    int n;
    vector<ll> sum;      // sum[node] = sum over the node's range, with the node's own tag applied
    vector<Affine> tag;  // tag[node] = map still waiting to be pushed to the two children
 
    explicit LazySeg(const vector<ll> &a) : n((int)a.size()), sum(4 * a.size() + 4, 0), tag(4 * a.size() + 4, IDENTITY) {
        build(1, 0, n, a);
    }
 
    void build(int node, int lo, int hi, const vector<ll> &a) {
        if (hi - lo == 1) {
            sum[node] = ((a[lo] % MOD) + MOD) % MOD;
            return;
        }
        int mid = (lo + hi) / 2;
        build(2 * node, lo, mid, a);
        build(2 * node + 1, mid, hi, a);
        sum[node] = (sum[2 * node] + sum[2 * node + 1]) % MOD;
    }
 
    // Apply a map to a whole node. The node's sum changes now; its children wait for the tag.
    // Every element of the node has the map applied once, so the sum of len values becomes a*sum + b*len.
    void apply(int node, int lo, int hi, Affine f) {
        ll len = hi - lo;
        sum[node] = (f.a * sum[node] + f.b * len) % MOD;
        tag[node] = then(tag[node], f);  // old pending map first, then the new one
    }
 
    // Move the node's pending tag down to its two children, then clear it.
    void push(int node, int lo, int hi) {
        int mid = (lo + hi) / 2;
        apply(2 * node, lo, mid, tag[node]);
        apply(2 * node + 1, mid, hi, tag[node]);
        tag[node] = IDENTITY;
    }
 
    void update(int node, int lo, int hi, int l, int r, Affine f) {
        if (r <= lo || hi <= l) return;  // no overlap
        if (l <= lo && hi <= r) {        // fully covered
            apply(node, lo, hi, f);
            return;
        }
        push(node, lo, hi);  // partial overlap: children must be up to date first
        int mid = (lo + hi) / 2;
        update(2 * node, lo, mid, l, r, f);
        update(2 * node + 1, mid, hi, l, r, f);
        sum[node] = (sum[2 * node] + sum[2 * node + 1]) % MOD;
    }
 
    ll query(int node, int lo, int hi, int l, int r) {
        if (r <= lo || hi <= l) return 0;
        if (l <= lo && hi <= r) return sum[node];
        push(node, lo, hi);
        int mid = (lo + hi) / 2;
        return (query(2 * node, lo, mid, l, r) + query(2 * node + 1, mid, hi, l, r)) % MOD;
    }
 
    // Public wrappers: apply x -> a*x + b to every index in [l, r).
    void add_affine(int l, int r, ll a, ll b) { update(1, 0, n, l, r, {a % MOD, b % MOD}); }
    // Sum of values in [l, r), mod MOD.
    ll range_sum(int l, int r) { return query(1, 0, n, l, r); }
};

The program that reads the problem is short. It builds the tree once and then handles each line in order.

#include "lazy_tree.h"
 
// Input: n q, then n values, then q lines.
//   "1 l r a b": every x in [l, r) becomes a*x + b (mod MOD)
//   "2 l r":     print the sum of x in [l, r) (mod MOD)
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, q;
    cin >> n >> q;
    vector<ll> a(n);
    for (auto &x : a) cin >> x;
    LazySeg t(a);
    while (q--) {
        int type, l, r;
        cin >> type >> l >> r;
        if (type == 1) {
            ll A, B;
            cin >> A >> B;
            t.add_affine(l, r, A, B);
        } else {
            cout << t.range_sum(l, r) << "\n";
        }
    }
    return 0;
}

In the header, push is called only when a node partly overlaps the query or update. A leaf either lies inside the range or outside it, so push never reaches below the leaves.

Worked original problem: the affine ledger

You have n cells, each holding a number. There are q operations. An operation either rewrites a range of cells with an affine map or asks for the sum of a range. Every result is taken modulo 998244353.

Input format. The first line has n and q. The second line has the n starting values, each up to 10^9. Each of the next q lines is one of these:

1 l r a b: for every index i in [l, r), replace the value x_i with a * x_i + b (mod 998244353). Here 0 <= a < 998244353 and 0 <= b < 998244353.
2 l r: print the sum of x_i for i in [l, r), modulo 998244353.

Limits: 1 <= n, q <= 2 * 10^5, and 0 <= l < r <= n.

Sample input:

Walk through it by hand first. The first update is 1 1 4 1 10. It maps every value in indices 1 to 3 with a = 1, b = 10. The array is now [1, 12, 13, 14, 5]. The second update is 1 2 5 0 7. It assigns 7 to indices 2 to 4, so the array becomes [1, 12, 7, 7, 7]. The first query sums the whole array, 1 + 12 + 7 + 7 + 7 = 34. The second query sums indices 1 and 2, 12 + 7 = 19.

Run the program on the sample:

34
19

The order demonstration uses the same array with the two updates in both orders. It is in the file order.cpp, which includes the same header:

add then assign: sum[0,5) = 34
assign then add: sum[0,5) = 54

Assigning first makes the array [1, 2, 7, 7, 7]. Adding 10 to indices 1 to 3 then gives [1, 12, 17, 17, 7], whose sum is 54. The two numbers differ, which confirms that the composition depends on the order, as the header's rule requires.

Stress test against a plain array

The stress program makes random arrays with up to 12 cells and random sequences of up to 40 operations. Each operation is either an affine update or a sum query. Updates are applied to the tree and to a plain array. Queries are compared with a plain loop. At the end of each round, every cell is compared too. The map is assign with probability about 1/4, add with probability about 1/4, and a random affine map otherwise.

all 2000 rounds passed, 33296 sum checks
all 20000 rounds passed, 334564 sum checks

The tree passed on these inputs. That shows the code agrees with the plain array on 20000 random sequences with a fixed seed. It does not prove the code is correct for every input.

A separate test checks the composition rule alone. It picks 20000 random pairs of maps and a random value, and it checks that applying the composed map equals applying the two maps one after the other. The check is in Python, so the site can run it. It has no input.

Python · editable

Its output on the machine that wrote this lesson:

checked 20000 random pairs, mismatches: 0

A bug the stress test catches

To see the check work, change one line in a copy of the header. Replace then(tag[node], f) with then(f, tag[node]). This applies the new map before the old one. The stress test, built against the copy, reports a mismatch on the first round where two non-trivial updates reach the same node:

MISMATCH round=1 op=11 l=2 r=5 got=560439675 want=288178111

A wrong sum like this is the typical symptom of a tag composition error. Nothing crashes, so only the comparison with the plain array shows the problem.

Complexity

Each update or query visits O(log n) nodes, because it follows at most two paths from the root and pushes at most two tags per level. Building takes O(n). The tree uses about 4n nodes of long long sums and affine tags.

On the machine that wrote this lesson, a random input with n = q = 2 * 10^5 ran in about 0.64 seconds of user time with g++ -O2. That time depends on the machine and the compiler, so treat it as an example and not a benchmark.

Common mistakes

Forgetting to push before recursing into a partial overlap. The children then hold old values, and a query reads a number from before the last update. The fix is the push call before the two recursive calls in update and query.
Composing the tags in the wrong order. The one-line change shown above produces a mismatch on the first round where two updates interact. Write the composition as "old first, new second" and test it with two maps by hand.
Forgetting that a sum depends on the segment length. An add of b to a node of length len changes the sum by b * len, not by b. Leaving out len gives correct answers only for leaves.
Pushing a tag that has already been cleared. Reset the tag to the identity map after pushing, not before. If you reset first, the children receive the identity and the update is lost.
Letting numbers grow without a modulus. With long long and no reduction, the product a * sum overflows after a few updates. Reduce after every multiplication when the problem asks for a modulus.
Mixing inclusive and half-open ranges. The header uses [l, r) everywhere. A call with r meant as the last index silently skips the last cell.

Checklist to keep

Write the invariant down before you code it: the node sum includes the node's own tag, and the tag is pending for the children.
Push before you read or change a child. Push only at partial overlaps.
Compose tags as old first, then new. Test with two maps by hand on a three-cell array.
Use the segment length in the sum update for any affine or add tag.
Keep a plain array next to the tree in tests, and compare every query and every cell at the end.
Reduce modulo the problem's modulus after each product.
Use one range convention, half-open, for the whole program.

Practice problems

Each answer sits under the line "Show the answer".

Compose by hand. Let f be x -> 2x + 1 and g be x -> 3x + 5. Find the single map (a, b) that equals g applied after f. Check it at x = 1.

Show the answer

g(f(x)) = 3 * (2x + 1) + 5 = 6x + 8, so (a, b) = (6, 8). At x = 1, f gives 3, g gives 14, and 6 * 1 + 8 = 14. The composition rule in the header gives (g.a * f.a, g.a * f.b + g.b) = (3 * 2, 3 * 1 + 5) = (6, 8).

Range add with range minimum. Explain how the node update changes when the tree stores the minimum of each segment instead of the sum. What is the new value of a node after adding v to its whole segment, and why is no change needed in the children until a push?

Show the answer

The minimum of a segment increases by v when every element increases by v, so the node's minimum becomes old minimum + v. The sum formula a * s + b * len is replaced by min + v. The tag still composes as "add v", and the children learn about the change only at a push, exactly as in the sum tree.

Assign then add. Take a node with an assign tag (0, 7) already pending, then receive "add 3" as (1, 3). Compute the composed tag with the header's rule and check it on the value 100.

Show the answer

then((0, 7), (1, 3)) = (1 * 0, 1 * 7 + 3) = (0, 10). On the value 100, the first map gives 7 and the second gives 10. The composed map gives 0 * 100 + 10 = 10. Both paths agree.

Find the bug. A colleague wrote this push function. Explain what goes wrong and fix it.

void push(int node, int lo, int hi) {
    int mid = (lo + hi) / 2;
    tag[node] = IDENTITY;
    apply(2 * node, lo, mid, tag[node]);
    apply(2 * node + 1, mid, hi, tag[node]);
}

Show the answer

The reset happens before the apply, so both children receive the identity map and the pending update is lost. Save the tag, apply it, then reset: Affine t = tag[node]; apply(2 * node, lo, mid, t); apply(2 * node + 1, mid, hi, t); tag[node] = IDENTITY;. The header's version passes tag[node] into apply before it resets the tag.

Write the naive checker. Write a function that takes the array as a vector and an operation list, applies every operation with a plain loop, and returns the answer to every sum query. Use it as the reference in your own test. Say what data type the sums need for the sample limits.

Show the answer

Apply each update with for i in [l, r): x[i] = (a * x[i] + b) % MOD and compute each sum with for i in [l, r): s += x[i], reducing modulo MOD. With n and q up to 2 * 10^5 the naive loop is too slow for the real limits, which is why it is only a reference for small tests. The reduced values are below 998244353, so each product of two values is below about 10^18 and fits in a signed 64-bit integer.

Exercise on your own machine

Do this in an empty folder, with g++ -O2 -std=c++17.

Save the header from this lesson as lazy_tree.h and the program that reads the problem as solve.cpp.
Save the sample input above as sample.txt. Compile with g++ -O2 -std=c++17 -o solve solve.cpp, then run ./solve < sample.txt. You should see 34 and 19.
Write stress.cpp that includes lazy_tree.h. Make it generate random arrays and operations as described above, and compare every answer with a plain array. Run it with 20000 rounds.
Copy the header to bug_tree.h and change the composition line. Build the stress test against the copy and check that it reports a mismatch.

Your numbers will differ from the ones in this lesson if you change the seed or the limits. Your timing will also differ, so do not compare it with the lesson's figure.

Summary

Lazy propagation keeps range updates at O(log n) by storing a pending map on each node and pushing it down only when a query or update reaches the children. The map is affine, so add, assign and multiply all fit in one form. Compose maps old first, then new. Test every tree against a plain array, because an error in the order of operations gives a wrong number and no crash.