Verify/RangeAffineRangeSum.test.cpp

• View this file on GitHub
• Last update: 2023-02-11 22:28:05+08:00
• Problem: https://judge.yosupo.jp/problem/range_affine_range_sum

Depends on

• DataStructure/LazySegmentTree.hpp
• ModInt/Modint32.hpp
• Template/Template.hpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/range_affine_range_sum"

#include "../DataStructure/LazySegmentTree.hpp"
#include "../ModInt/Modint32.hpp"
#include "../Template/Template.hpp"

constexpr int mod = 998244353;
using Z = mint<mod>;
struct Tag {
    Z x = 1;
    Z y = 0;
    void apply(const Tag& a) {
        x = x * a.x;
        y = a.x * y + a.y;
    }
};
struct Info {
    Z x = 0;
    int l, r;
    void apply(const Tag& a) {
        x *= a.x;
        x += Z(r - l) * a.y;
    }
};
Info operator+(const Info& a, const Info& b) { return {a.x + b.x, a.l, b.r}; }
int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;

    std::vector<Info> a(n);
    for (int i = 0; i < n; i++) {
        std::cin >> a[i].x;
        a[i].l = i;
        a[i].r = i + 1;
    }

    LazySegmentTree<Info, Tag> t(a);

    for (int i = 1; i <= q; i++) {
        int op;
        std::cin >> op;
        if (op == 0) {
            int l, r, b, c;
            std::cin >> l >> r >> b >> c;
            t.rangeApply(l, r, Tag{b, c});
        }
        if (op == 1) {
            int l, r;
            std::cin >> l >> r;
            auto ans = t.rangeQuery(l, r);
            std::cout << ans.x << '\n';
        }
    }

    return 0;
}

#line 1 "Verify/RangeAffineRangeSum.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/range_affine_range_sum"

#line 2 "Template/Template.hpp"

#include <bits/stdc++.h>

using i64 = std::int64_t;
#line 3 "DataStructure/LazySegmentTree.hpp"
template <class Info, class Tag>
struct LazySegmentTree {
    /*
        Info: apply,operator +
        Tag: apply
        all operations obey [l,r)
    */
    const int n;
    std::vector<Info> info;
    std::vector<Tag> tag;
    LazySegmentTree(int n)
        : n(n), info(4 << std::__lg(n)), tag(4 << std::__lg(n)) {}
    LazySegmentTree(std::vector<Info> init) : LazySegmentTree(init.size()) {
        std::function<void(int, int, int)> build = [&](int p, int l, int r) {
            if (r - l == 1) {
                info[p] = init[l];
                return;
            }
            int m = (l + r) / 2;
            build(2 * p, l, m);
            build(2 * p + 1, m, r);
            pull(p);
        };
        build(1, 0, n);
    }
    void pull(int p) { info[p] = info[2 * p] + info[2 * p + 1]; }
    void apply(int p, const Tag &v) {
        info[p].apply(v);
        tag[p].apply(v);
    }
    void push(int p) {
        apply(2 * p, tag[p]);
        apply(2 * p + 1, tag[p]);
        tag[p] = Tag();
    }
    void modify(int p, int l, int r, int x, const Info &v) {
        if (r - l == 1) {
            info[p] = v;
            return;
        }
        int m = (l + r) / 2;
        push(p);
        if (x < m) {
            modify(2 * p, l, m, x, v);
        } else {
            modify(2 * p + 1, m, r, x, v);
        }
        pull(p);
    }
    void modify(int p, const Info &v) { modify(1, 0, n, p, v); }
    Info rangeQuery(int p, int l, int r, int x, int y) {
        if (l >= y || r <= x) {
            return Info();
        }
        if (l >= x && r <= y) {
            return info[p];
        }
        int m = (l + r) / 2;
        push(p);
        return rangeQuery(2 * p, l, m, x, y) +
               rangeQuery(2 * p + 1, m, r, x, y);
    }
    Info rangeQuery(int l, int r) { return rangeQuery(1, 0, n, l, r); }
    void rangeApply(int p, int l, int r, int x, int y, const Tag &v) {
        if (l >= y || r <= x) {
            return;
        }
        if (l >= x && r <= y) {
            apply(p, v);
            return;
        }
        int m = (l + r) / 2;
        push(p);
        rangeApply(2 * p, l, m, x, y, v);
        rangeApply(2 * p + 1, m, r, x, y, v);
        pull(p);
    }
    void rangeApply(int l, int r, const Tag &v) {
        return rangeApply(1, 0, n, l, r, v);
    }
};
#line 2 "ModInt/Modint32.hpp"

#line 4 "ModInt/Modint32.hpp"

template <int mod>
struct mint {
    int x;
    mint() : x(0) {}
    mint(int64_t y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}
    mint &operator+=(const mint &p) {
        if ((x += p.x) >= mod) x -= mod;
        return *this;
    }
    mint &operator-=(const mint &p) {
        if ((x += mod - p.x) >= mod) x -= mod;
        return *this;
    }
    mint &operator*=(const mint &p) {
        x = (int)(1LL * x * p.x % mod);
        return *this;
    }
    mint &operator/=(const mint &p) {
        *this *= p.inverse();
        return *this;
    }
    mint operator-() const { return mint(-x); }
    mint operator+(const mint &p) const { return mint(*this) += p; }
    mint operator-(const mint &p) const { return mint(*this) -= p; }
    mint operator*(const mint &p) const { return mint(*this) *= p; }
    mint operator/(const mint &p) const { return mint(*this) /= p; }
    bool operator==(const mint &p) const { return x == p.x; }
    bool operator!=(const mint &p) const { return x != p.x; }
    mint inverse() const {
        int a = x, b = mod, u = 1, v = 0, t;
        while (b > 0) {
            t = a / b;
            std::swap(a -= t * b, b);
            std::swap(u -= t * v, v);
        }
        return mint(u);
    }
    friend std::ostream &operator<<(std::ostream &os, const mint &p) {
        return os << p.x;
    }
    friend std::istream &operator>>(std::istream &is, mint &a) {
        int64_t t;
        is >> t;
        a = mint<mod>(t);
        return (is);
    }
    int get() const { return x; }
    static constexpr int get_mod() { return mod; }
};
#line 6 "Verify/RangeAffineRangeSum.test.cpp"

constexpr int mod = 998244353;
using Z = mint<mod>;
struct Tag {
    Z x = 1;
    Z y = 0;
    void apply(const Tag& a) {
        x = x * a.x;
        y = a.x * y + a.y;
    }
};
struct Info {
    Z x = 0;
    int l, r;
    void apply(const Tag& a) {
        x *= a.x;
        x += Z(r - l) * a.y;
    }
};
Info operator+(const Info& a, const Info& b) { return {a.x + b.x, a.l, b.r}; }
int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;

    std::vector<Info> a(n);
    for (int i = 0; i < n; i++) {
        std::cin >> a[i].x;
        a[i].l = i;
        a[i].r = i + 1;
    }

    LazySegmentTree<Info, Tag> t(a);

    for (int i = 1; i <= q; i++) {
        int op;
        std::cin >> op;
        if (op == 0) {
            int l, r, b, c;
            std::cin >> l >> r >> b >> c;
            t.rangeApply(l, r, Tag{b, c});
        }
        if (op == 1) {
            int l, r;
            std::cin >> l >> r;
            auto ans = t.rangeQuery(l, r);
            std::cout << ans.x << '\n';
        }
    }

    return 0;
}

Back to top page