🗂️Data StructureAdvanced

Segment Tree with Range Affine Transformation

Key Points

  • •
    A segment tree with lazy propagation can support range updates of the form + b (affine transformations) and range-sum queries in O(log n) per operation.
  • •
    The lazy tag is a pair (a, b) representing the transformation f(x) = a·x + b; identity is (1, 0) and assignment, addition, multiplication become special cases.
  • •
    When a node covering len elements with current sum S receives tag (a, b), its new sum becomes a·S + b·len, because every element is transformed the same way.
  • •
    Tags compose as (a2, b2) ∘ (a1, b1) = (a2·a1, a2·b1 + b2), meaning apply (a1, b1) first, then (a2, b2).
  • •
    Lazy propagation stores pending tags at internal nodes and pushes them to children only when necessary, keeping updates fast.
  • •
    Using modulo arithmetic (e.g., 998244353) is common to prevent overflow and to match competitive programming problem statements.
  • •
    This structure is ideal when you need to mix range assignment, range addition, and range multiplication together with range-sum queries.
  • •
    Be careful with the order of tag composition, handling the segment length in b·len, and resetting/identifying the identity tag.

Prerequisites

  • →Segment Tree Basics — Understanding how segment trees partition ranges and how to implement push/pull is essential.
  • →Function Composition — Lazy tags are affine functions that must be composed in the correct order.
  • →Modular Arithmetic — Most problems require modulo operations to avoid overflow and match problem statements.
  • →Asymptotic Analysis — To reason why each operation is O(log n) and the build is O(n).
  • →Recursion and Tree Traversal — Updates and queries are implemented recursively over the tree.

Detailed Explanation

Tap terms for definitions

01Overview

A segment tree with range affine transformation is a balanced binary tree data structure that supports two core operations: updating every element within an interval [l, r] by the same affine function f(x) = a·x + b, and querying the sum of values over an interval. Affine means a linear scaling (multiply by a) followed by a shift (add b). The key idea is that if you apply this same function to every element in a segment, the segment’s total sum transforms predictably: the sum S becomes a·S + b·len, where len is the number of elements in the segment. To achieve logarithmic time, the tree uses lazy propagation: instead of updating each element immediately, it stores a “pending operation” (a lazy tag) at nodes that fully lie in the update range and pushes that operation down to children only when those children are visited later. The lazy tag is exactly the affine pair (a, b), with identity (1, 0). Tags must be composed correctly because multiple updates can stack. Composition is function composition: applying (a2, b2) after (a1, b1) equals (a2·a1, a2·b1 + b2). With this design, both updates and queries run in O(log n), build runs in O(n), and memory is O(n). This pattern is widely used in competitive programming and can be adapted for modular arithmetic, which is common to avoid overflow and conform to problem constraints.

02Intuition & Analogies

Imagine a long fence with n boards. You frequently perform bulk operations on contiguous stretches of boards: multiply the colors’ brightness by a, then add b to adjust the shade. This is like taking each board’s color value x and replacing it with a·x + b. If you just want to know the total brightness of a section, you don’t need the exact color of each individual board; you only need the total sum of that section. Now, if every board’s color in a section is multiplied by a, the total brightness of that section also multiplies by a. If you add b to every board, the total increases by b times the number of boards in the section. That’s why the new sum is a·S + b·len. The segment tree is like organizing the fence into a hierarchy of intervals: the whole fence at the root, halves below, quarters below that, and so on. When you decide to recolor a whole interval (say, boards 100–200), you slap a sticky note on the node representing that interval that says “apply (a, b) later.” You also immediately adjust that node’s stored total sum using the simple formula above. You don’t propagate to children right away; you wait until you actually need to look inside them for a subquery. When you eventually go deeper, you peel off the sticky note and hand copies to the children, updating their sums consistently. This ‘lazy’ strategy saves time because many queries and updates only touch O(log n) nodes rather than all elements. Finally, if you apply multiple sticky notes over time, they stack like function composition: stretch-then-shift, then maybe another stretch-then-shift, and so on; the order matters, and composing them correctly is the heart of the method.

03Formal Definition

Let A[1..n] be an array over a ring (e.g., integers modulo M). We support: (1) Range update on [l, r]: for all i in [l, r], set + b, with parameters a, b from the ring; (2) Range sum query on [l, r]: return S = A[i]. We maintain a segment tree where each node v covers an interval and stores = A[i] and le = . Each node also stores a lazy tag = (, ) representing an affine transformation to be applied to all elements in (and to be pushed to its children when needed). The identity tag is (1, 0). Applying tag T = (a, b) to node v updates ← a· + b·le and composes the node’s lazy tag as ← T ∘ , where composition is defined as (a2, b2) ∘ (a1, b1) = (a2·a1, a2·b1 + b2). When traversing from a node to its children, we push to the children: for each child u, apply to and update ← ∘ ; then reset to (1, 0). Queries and updates recurse over the tree, splitting at most O(log n) times because the tree height is O(log n). The structure works over any commutative ring supporting + and ·, commonly using modular arithmetic to prevent overflow.

04When to Use

Use this structure when you need to combine range-sum queries with multiple kinds of range updates that can all be expressed as x → a·x + b. This includes range assignment (set to v) via (a, b) = (0, v), range addition via (1, v), range multiplication via (v, 0), and any composition of these in arbitrary order. It is especially suitable for problems that involve sequences of linear transformations, simulating repeated operations, or applying linear recurrences where each step is of the form a·x + b across a whole interval. Competitive programming platforms (AtCoder, Codeforces) often feature tasks requiring mixing assignment and add/multiply operations under modulo. A Fenwick tree (BIT) is great for range add/point query or point update/range sum, but it cannot handle assignment and multiplication together; even range add + range multiply together on sums is tricky for BIT without more structure. A naive per-element update is O(n) per operation and too slow for large constraints, whereas this segment tree keeps both updates and queries to O(log n). If your queries are only point updates and range sums, a Fenwick tree is simpler. If your updates are only range add (no multiply/assign), a simpler lazy segment tree with additive tags also suffices. Use this affine-lazy segment tree when you truly need the full flexibility of linear (affine) range updates.

⚠️Common Mistakes

• Wrong composition order. If T_old is already pending and a new T_new arrives, the correct stacking is T_old ← T_new ∘ T_old because the older transform must be applied first at the leaves. Reversing the order silently produces wrong answers. • Forgetting the length in the sum update. When applying (a, b) to a node sum S over len elements, the updated sum is a·S + b·len, not a·S + b. • Not normalizing modulo. Under modulo arithmetic, reduce a, b, and sums after every operation to keep values in range. Negative values must be handled with (x % MOD + MOD) % MOD. • Forgetting to push before descending. During a partial-overlap query or update, push the lazy tag to children before recursing; otherwise, child sums and tags become inconsistent. • Confusing identity with zero. The identity affine tag is (1, 0). Using (0, 0) as an ‘empty’ tag mistakenly wipes values during composition. • Overflow in 64-bit integers. Products a·S and b·len can exceed 64-bit; prefer modular arithmetic or use wider intermediates (e.g., __int128) and beware that repeated composition can still overflow. • Not storing the segment length. You need len to compute b·len in constant time; recomputing length from indices at every node is error-prone and slow.

Key Formulas

Identity

Explanation: Applying (1, 0) leaves every x unchanged: 1·x + 0 = x. Composing with identity does not change other tags.

Tag Composition

Explanation: Apply (, ) first, then (, ). This preserves the order of updates in time.

Sum Update Under Affine

Explanation: If every element x in a segment is replaced by a·x + b, then the total sum increases linearly by scaling S and adding b times the number of elements.

Assignment as Tag

Explanation: Setting x to v is equivalent to 0·x + v for any x. This overwrites previous values.

Addition as Tag

Explanation: Adding v to every x corresponds to an affine tag with a = 1 and b = v.

Multiplication as Tag

Explanation: Scaling every x by v corresponds to an affine tag with a = v and b = 0.

Repeated Affine Application

Explanation: Applying the same affine function k times yields a geometric series in a for the additive term. Useful for reasoning about many repeated updates.

Matrix Form of Affine

Explanation: Affine transformations can be represented as 2×2 matrices; composition becomes matrix multiplication, which explains the composition rule.

Sum Update via Matrix

Explanation: Treat the segment sum and length as a vector. Applying the affine matrix shows directly that S' = a·S + b·len while len stays unchanged.

Per-Operation Cost

Explanation: Each update or query visits at most one path down the tree and splits into two at most O(log n) levels, yielding logarithmic time.

Complexity Analysis

Let n be the array length. The segment tree has height O(log n) because it repeatedly halves intervals. A range update or range query descends from the root, and at each level at most two nodes are visited due to interval splitting; thus the number of visited nodes per operation is O(log n). Each visited node performs O(1) work: applying a tag (constant-time arithmetic), composing tags (constant-time arithmetic), and a few assignments. Therefore the time per operation is O(log n). Building the tree from an array takes O(n) time by recursively combining children or iteratively initializing leaves and pulling upward. Space usage is O(n): we allocate arrays of size about 4n to store nodes, sums, lengths, and lazy tags safely for any n (a constant factor overhead). Constants matter: each apply involves two multiplications and one addition to update the sum for S' = a·S + b·len, and a tag composition requires one multiplication and one addition for b (plus one multiplication for a). With modular arithmetic, each operation adds one or two reductions. In practice, this is fast enough for constraints up to ~2e5–1e6 operations. Edge cases affecting performance include frequent partial-overlap updates forcing pushes at many levels; nonetheless, the number of pushes is bounded by O(log n) per operation. There is no amortization trick needed: the worst-case time per operation is already O(log n).

Code Examples

Segment Tree with Affine Lazy Propagation (Modulo 998244353)
1#include <bits/stdc++.h>
2using namespace std;
3
4// Segment tree supporting range affine updates x -> a*x + b and range sum queries, modulo MOD.
5static const long long MOD = 998244353LL;
6
7long long norm(long long x) {
8    x %= MOD;
9    if (x < 0) x += MOD;
10    return x;
11}
12
13struct Tag { // affine transform f(x) = a*x + b
14    long long a = 1; // multiplier
15    long long b = 0; // adder
16};
17
18struct Node {
19    long long sum = 0; // sum over the interval (mod MOD)
20    int len = 0;       // length of the interval
21    Tag lazy;          // pending affine transform
22};
23
24struct SegTree {
25    int n;                 // size of the input array
26    vector<Node> st;       // segment tree nodes (1-indexed internally)
27
28    SegTree(const vector<long long>& arr) { build(arr); }
29
30    static Tag compose(const Tag& f, const Tag& g) {
31        // Return f ∘ g: apply g first, then f.
32        // (a2,b2) ∘ (a1,b1) = (a2*a1, a2*b1 + b2)
33        Tag h;
34        h.a = (f.a * g.a) % MOD;
35        h.b = ( (f.a * g.b) + f.b ) % MOD;
36        return h;
37    }
38
39    static void apply_tag(Node& nd, const Tag& t) {
40        // Update the node's sum by the affine transform and compose its lazy tag.
41        nd.sum = ( (t.a * nd.sum) + (t.b * nd.len) ) % MOD;
42        nd.lazy = compose(t, nd.lazy);
43    }
44
45    void push(int p) {
46        // Push lazy tag at node p to its children (if any)
47        Tag &lz = st[p].lazy;
48        if (lz.a != 1 || lz.b != 0) {
49            apply_tag(st[p<<1], lz);
50            apply_tag(st[p<<1|1], lz);
51            lz = Tag{}; // reset to identity (1,0)
52        }
53    }
54
55    void pull(int p) {
56        st[p].sum = (st[p<<1].sum + st[p<<1|1].sum) % MOD;
57    }
58
59    void build(const vector<long long>& arr) {
60        n = (int)arr.size();
61        st.assign(4*n + 5, Node{});
62        build(1, 0, n-1, arr);
63    }
64
65    void build(int p, int l, int r, const vector<long long>& arr) {
66        st[p].len = r - l + 1;
67        st[p].lazy = Tag{}; // identity
68        if (l == r) {
69            st[p].sum = norm(arr[l]);
70            return;
71        }
72        int m = (l + r) >> 1;
73        build(p<<1, l, m, arr);
74        build(p<<1|1, m+1, r, arr);
75        pull(p);
76    }
77
78    // Apply affine transform t to interval [L, R]
79    void range_apply(int L, int R, Tag t) {
80        t.a = norm(t.a);
81        t.b = norm(t.b);
82        range_apply(1, 0, n-1, L, R, t);
83    }
84
85    void range_apply(int p, int l, int r, int L, int R, const Tag& t) {
86        if (R < l || r < L) return;
87        if (L <= l && r <= R) { apply_tag(st[p], t); return; }
88        push(p);
89        int m = (l + r) >> 1;
90        range_apply(p<<1, l, m, L, R, t);
91        range_apply(p<<1|1, m+1, r, L, R, t);
92        pull(p);
93    }
94
95    // Query sum over [L, R]
96    long long range_query(int L, int R) { return range_query(1, 0, n-1, L, R); }
97
98    long long range_query(int p, int l, int r, int L, int R) {
99        if (R < l || r < L) return 0;
100        if (L <= l && r <= R) return st[p].sum;
101        push(p);
102        int m = (l + r) >> 1;
103        long long left = range_query(p<<1, l, m, L, R);
104        long long right = range_query(p<<1|1, m+1, r, L, R);
105        return (left + right) % MOD;
106    }
107};
108
109int main() {
110    ios::sync_with_stdio(false);
111    cin.tie(nullptr);
112
113    // Example usage
114    int n = 8;
115    vector<long long> a = {1,2,3,4,5,6,7,8};
116    SegTree st(a);
117
118    // Range addition: add 5 to [2,6]
119    st.range_apply(2, 6, Tag{1, 5});
120    // Range multiplication: multiply by 3 on [0,3]
121    st.range_apply(0, 3, Tag{3, 0});
122    // Range assignment: set [4,7] to value 10
123    st.range_apply(4, 7, Tag{0, 10});
124
125    // Query some sums
126    cout << st.range_query(0, 7) << "\n"; // total sum
127    cout << st.range_query(2, 5) << "\n"; // subrange sum
128
129    return 0;
130}
131

This implementation maintains for each node the sum modulo MOD and the segment length. The lazy tag (a, b) is composed in the correct temporal order: new_tag ∘ old_tag. Applying a tag updates the node’s sum via S' = a·S + b·len and stores the composed tag for later propagation. Range assignment, addition, and multiplication become special cases: (0, v), (1, v), and (v, 0). The main function demonstrates a sequence of mixed updates and queries.

Time: O(log n) per update or query; O(n) buildSpace: O(n)
64-bit (non-mod) Variant Using __int128 Intermediates
1#include <bits/stdc++.h>
2using namespace std;
3
4// WARNING: Without modulo, values can overflow 64-bit if a, b, or sums grow large.
5// We use __int128 for intermediates to reduce accidental overflow,
6// but this does not prevent true overflows of long long results.
7
8struct Tag { long long a = 1, b = 0; };
9struct Node {
10    long long sum = 0; // sum over the interval
11    int len = 0;       // segment length
12    Tag lazy;          // pending affine tag
13};
14
15struct SegTree {
16    int n; vector<Node> st;
17    SegTree(const vector<long long>& arr) { build(arr); }
18
19    static Tag compose(const Tag& f, const Tag& g) {
20        // f ∘ g
21        __int128 a = (__int128)f.a * g.a;
22        __int128 b = (__int128)f.a * g.b + f.b;
23        Tag h; h.a = (long long)a; h.b = (long long)b; return h;
24    }
25    static void apply_tag(Node& nd, const Tag& t) {
26        __int128 ns = (__int128)t.a * nd.sum + (__int128)t.b * nd.len;
27        nd.sum = (long long)ns; // may overflow if ns exceeds 64-bit range
28        nd.lazy = compose(t, nd.lazy);
29    }
30    void push(int p) {
31        if (st[p].lazy.a != 1 || st[p].lazy.b != 0) {
32            apply_tag(st[p<<1], st[p].lazy);
33            apply_tag(st[p<<1|1], st[p].lazy);
34            st[p].lazy = Tag{};
35        }
36    }
37    void pull(int p) { st[p].sum = st[p<<1].sum + st[p<<1|1].sum; }
38
39    void build(const vector<long long>& arr) {
40        n = (int)arr.size(); st.assign(4*n+5, Node{}); build(1,0,n-1,arr);
41    }
42    void build(int p, int l, int r, const vector<long long>& arr) {
43        st[p].len = r-l+1; st[p].lazy = Tag{};
44        if (l==r) { st[p].sum = arr[l]; return; }
45        int m=(l+r)>>1; build(p<<1,l,m,arr); build(p<<1|1,m+1,r,arr); pull(p);
46    }
47
48    void range_apply(int L, int R, Tag t){ range_apply(1,0,n-1,L,R,t);}    
49    void range_apply(int p,int l,int r,int L,int R,const Tag& t){
50        if(R<l||r<L) return; if(L<=l&&r<=R){ apply_tag(st[p],t); return; }
51        push(p); int m=(l+r)>>1; range_apply(p<<1,l,m,L,R,t); range_apply(p<<1|1,m+1,r,L,R,t); pull(p);
52    }
53
54    long long range_query(int L,int R){ return range_query(1,0,n-1,L,R);}    
55    long long range_query(int p,int l,int r,int L,int R){
56        if(R<l||r<L) return 0; if(L<=l&&r<=R) return st[p].sum;
57        push(p); int m=(l+r)>>1; return range_query(p<<1,l,m,L,R)+range_query(p<<1|1,m+1,r,L,R);
58    }
59};
60
61int main(){
62    ios::sync_with_stdio(false); cin.tie(nullptr);
63
64    vector<long long> a = {10, 20, 30, 40, 50};
65    SegTree st(a);
66
67    // Add 5 to [1,3]
68    st.range_apply(1,3, Tag{1,5});
69    // Multiply by 2 on [0,4]
70    st.range_apply(0,4, Tag{2,0});
71    // Assign value 7 to [2,4]
72    st.range_apply(2,4, Tag{0,7});
73
74    cout << st.range_query(0,4) << "\n"; // sum of all elements
75    cout << st.range_query(2,4) << "\n"; // sum over [2,4]
76    return 0;
77}
78

Same API as the modular version, but without modulo arithmetic. This demonstrates how to implement and use affine-lazy segment trees over 64-bit integers. It uses __int128 for intermediate products and sums while computing a·S + b·len and composing tags to reduce accidental overflow during computation, but final results are cast back to long long. Use with caution if values can exceed 64-bit ranges, or prefer a modular implementation.

Time: O(log n) per operation; O(n) buildSpace: O(n)