Fading Coder

One Final Commit for the Last Sprint

Home > Tech > Content

Segment Tree with Composite Lazy Tags: Handling Range Add, Multiply and Assign Operations

Tech May 7 3

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 tot and element count len (needed for range addition)
  • Tag: composite transfromation parameters mul and add representing $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;
}

Related Articles

Understanding Strong and Weak References in Java

Strong References Strong reference are the most prevalent type of object referencing in Java. When an object has a strong reference pointing to it, the garbage collector will not reclaim its memory. F...

Comprehensive Guide to SSTI Explained with Payload Bypass Techniques

Introduction Server-Side Template Injection (SSTI) is a vulnerability in web applications where user input is improper handled within the template engine and executed on the server. This exploit can r...

Implement Image Upload Functionality for Django Integrated TinyMCE Editor

Django’s Admin panel is highly user-friendly, and pairing it with TinyMCE, an effective rich text editor, simplifies content management significantly. Combining the two is particular useful for bloggi...

Leave a Comment

Anonymous

◎Feel free to join the discussion and share your thoughts.