Segment Tree with Composite Lazy Tags: Handling Range Add, Multiply and Assign Operations
Given an array of $n$ integers $a_1, a_2, \ldots, a_n$, process $q$ range operations:
- Type 1
l r d: add $d$ to all elements in $[l, r]$ - Type 2
l r d: multiply all elements in $[l, r]$ by $d$ - Type 3
l r d: assign all elements in $[l, r]$ to $d$ - Type 4
l r: query the sum of elements in $[l, r]$ modulo $10^9+7$
Analysis
Each operation can be viewed as a linear transformation on array values. The assignment opertaion is equivalent to multiplying by zero then adding $d$. Thus all three operations fit the form $a \to a \times m + c$.
Node Structure
Maintain two pieces of information:
- Info: range sum
totand element countlen(needed for range addition) - Tag: composite transfromation parameters
mulandaddrepresenting $val \times mul + add$
Merging Rules
Info + Info: Simple addition of sums, addition of lengths.
$$res.tot = left.tot + right.tot, \quad res.len = left.len + right.len$$
Info transforemd by Tag: Apply the linear transformation to the sum.
$$new.tot = f.tot \times t.mul + f.len \times t.add, \quad new.len = f.len$$
Tag composed with Tag: Order matters! If tag $t_1$ applies before $t_2$:
$$(a \times t_1.mul + t_1.add) \times t_2.mul + t_2.add = a \times (t_1.mul \times t_2.mul) + (t_1.add \times t_2.mul + t_2.add)$$
Hence:
$$res.mul = t_1.mul \times t_2.mul, \quad res.add = t_1.add \times t_2.mul + t_2.add$$
Implementation Details
Initialize tags as identity: ${mul=1, add=0}$ since $a = a \times 1 + 0$.
#include <bits/stdc++.h>
using i64 = long long;
const i64 MOD = 1000000007;
const int MAXN = 200005;
int init[MAXN];
struct Data {
i64 tot, len;
Data() : tot(0), len(0) {}
Data(i64 _tot, i64 _len) : tot(_tot), len(_len) {}
void init(int idx) { tot = init[idx], len = 1; }
};
struct Lazy {
i64 mul, add;
Lazy() : mul(1), add(0) {}
Lazy(i64 _mul, i64 _add) : mul(_mul), add(_add) {}
};
Data mergeData(const Data& x, const Data& y) {
return Data((x.tot + y.tot) % MOD, x.len + y.len);
}
Data applyTag(const Data& d, const Lazy& z) {
return Data((d.tot * z.mul % MOD + d.len * z.add % MOD) % MOD, d.len);
}
Lazy compose(const Lazy& first, const Lazy& second) {
// first applies, then second applies
return Lazy(first.mul * second.mul % MOD,
(first.add * second.mul + second.add) % MOD);
}
struct SegNode {
Data val;
Lazy tag;
} tree[MAXN << 2];
void pull(int p) {
tree[p].val = mergeData(tree[p << 1].val, tree[p << 1 | 1].val);
}
void apply(int p, const Lazy& z) {
tree[p].val = applyTag(tree[p].val, z);
tree[p].tag = compose(tree[p].tag, z);
}
void push(int p) {
if (tree[p].tag.mul != 1 || tree[p].tag.add != 0) {
apply(p << 1, tree[p].tag);
apply(p << 1 | 1, tree[p].tag);
tree[p].tag = Lazy();
}
}
void build(int p, int l, int r) {
if (l == r) {
tree[p].val.init(l);
return;
}
int mid = (l + r) >> 1;
build(p << 1, l, mid);
build(p << 1 | 1, mid + 1, r);
pull(p);
}
Data query(int p, int l, int r, int ql, int qr) {
if (ql <= l && r <= qr) return tree[p].val;
push(p);
int mid = (l + r) >> 1;
if (qr <= mid) return query(p << 1, l, mid, ql, qr);
if (ql > mid) return query(p << 1 | 1, mid + 1, r, ql, qr);
return mergeData(query(p << 1, l, mid, ql, mid),
query(p << 1 | 1, mid + 1, r, mid + 1, qr));
}
void modify(int p, int l, int r, int ql, int qr, const Lazy& z) {
if (ql <= l && r <= qr) {
apply(p, z);
return;
}
push(p);
int mid = (l + r) >> 1;
if (ql <= mid) modify(p << 1, l, mid, ql, qr, z);
if (qr > mid) modify(p << 1 | 1, mid + 1, r, ql, qr, z);
pull(p);
}
int main() {
std::ios::sync_with_stdio(false);
std::cin.tie(nullptr);
int n, q;
std::cin >> n >> q;
for (int i = 1; i <= n; i++) std::cin >> init[i];
build(1, 1, n);
while (q--) {
int op, l, r, d;
std::cin >> op;
if (op == 1) {
std::cin >> l >> r >> d;
modify(1, 1, n, l, r, Lazy(1, d)); // add: mul=1, add=d
} else if (op == 2) {
std::cin >> l >> r >> d;
modify(1, 1, n, l, r, Lazy(d, 0)); // multiply: mul=d, add=0
} else if (op == 3) {
std::cin >> l >> r >> d;
modify(1, 1, n, l, r, Lazy(0, d)); // assign: mul=0, add=d
} else {
std::cin >> l >> r;
std::cout << query(1, 1, n, l, r).tot << '\n';
}
}
return 0;
}