test/yosupo/ds/segment_tree/range_affine_range_sum.test.cpp

• View this file on GitHub
• Last update: 2025-09-24 20:26:11+07:00
• Problem: https://judge.yosupo.jp/problem/range_affine_range_sum

Depends on

• Lazy Segment Tree (ds/lazysegtree.hpp)
• Modint (math/modint.hpp)
• other/update_proxy.hpp

Code

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

#include <bits/stdc++.h>
using namespace std;

#include "ds/lazysegtree.hpp"
#include "math/modint.hpp"

using mint = modint998;
struct M {
    using T = array<mint, 2>;
    using F = array<mint, 2>;
    static T id() { return {0, 0}; }
    static F fid() { return {1, 0}; }
    static T op(T l, T r) { return {l[0] + r[0], l[1] + r[1]}; }
    static F comp(F l, F r) { return {l[0] * r[0], l[0] * r[1] + l[1]}; }
    static bool map(F f, T &x) {
        x[0] *= f[0];
        x[0] += f[1] * x[1];
        return 1;
    }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    LazySegTree<M> st(n, [](int x) { return cin >> x, M::T{x, 1}; });
    while (q--) {
        int t, l, r, b, c;
        cin >> t >> l >> r;
        if (t == 0) {
            cin >> b >> c;
            st.apply(l, r, {b, c});
        } else {
            cout << st(l, r)[0] << "\n";
        }
    }
}

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

#include <bits/stdc++.h>
using namespace std;

#line 2 "ds/lazysegtree.hpp"

#line 2 "other/update_proxy.hpp"

// clang-format off
template<class T, class Cb> struct UpdateProxy {
    T &x;
    Cb cb;

    UpdateProxy(T &_x, Cb _cb) : x(_x), cb(_cb) {}

    operator const T &() const { return x; }
    auto &operator*() && { return x; }
    auto operator->() && { return &x; }
    auto operator++(int) && { T old = x; return ++x, cb(), old; }
    auto operator--(int) && { T old = x; return --x, cb(), old; }
    auto &operator++() && { ++x, cb(); return *this; }
    auto &operator--() && { --x, cb(); return *this; }
    auto &operator+=(const T &r) && { x += r, cb(); return *this; }
    auto &operator-=(const T &r) && { x -= r, cb(); return *this; }
    auto &operator*=(const T &r) && { x *= r, cb(); return *this; }
    auto &operator/=(const T &r) && { x /= r, cb(); return *this; }
    auto &operator%=(const T &r) && { x %= r, cb(); return *this; }
    auto &operator=(const T &r) && { x = r, cb(); return *this; }
    auto &operator<<=(const T &r) && { x <<= r, cb(); return *this; }
    auto &operator>>=(const T &r) && { x >>= r, cb(); return *this; }
    template<class F> auto &apply(F &&f) && { f(x), cb(); return *this; }
};
// clang-format on
#line 4 "ds/lazysegtree.hpp"

// M: T M::id(), F M::fid(), T op(T, T), F comp(F, F), bool map(F, &T)
template<class M> struct LazySegTree {
    using T = typename M::T;
    using F = typename M::F;

    int n, lg, m;
    vector<T> t;
    vector<char> upd;
    vector<F> lz;

    LazySegTree() = default;
    LazySegTree(int _n) :
        n(_n), lg(_lg(n)), m(1 << lg), t(2 * m, M::id()), upd(m),
        lz(m, M::fid()) {}
    template<class G> LazySegTree(int _n, G &&gen) : LazySegTree(_n) {
        for (int i = 0; i < n; i++) t[i + m] = gen(i);
        for (int i = m; --i;) update(i);
    }
    template<class V>
    LazySegTree(const V &v) :
        LazySegTree((int)v.size(), [&](int i) { return T(v[i]); }) {}

    vector<T> to_vec() {
        vector<T> r(n);
        for (int i = 0; i < n; i++) r[i] = (*this)[i];
        return r;
    }
    void set(int p, const T &x) {
        assert(0 <= p && p < n);
        push_to(p), t[p + m] = x, update_from(p);
    }
    void mul(int p, const T &x) {
        assert(0 <= p && p < n);
        push_to(p), t[p + m] = M::op(t[p + m], x), update_from(p);
    }
    auto operator[](int p) {
        assert(0 <= p && p < n);
        UpdateProxy up(t[p + m], [this, p]() { update_from(p); });
        return push_to(p), up;
    }
    T get(int p) { return (*this)[p]; }

    T operator()(int l, int r) {
        assert(0 <= l && l <= r && r <= n);
        if (l == r) return M::id();
        int li = __builtin_ctz(l + m), ri = __builtin_ctz(r + m);
        push_to(l, li), push_to(r - 1, ri);
        T sml = M::id(), smr = M::id();
        for (l += m, r += m; l < r; l >>= 1, r >>= 1) {
            if (l & 1) sml = M::op(sml, t[l++]);
            if (r & 1) smr = M::op(t[--r], smr);
        }
        return M::op(sml, smr);
    }

    T pref(int r) { return (*this)(0, r); }
    T suff(int l) { return (*this)(l, n); }
    T prod(int l, int r) { return (*this)(l, r); }
    T all_prod() const { return t[1]; }

    void apply(int p, const F &f) {
        assert(0 <= p && p < n);
        push_to(p), M::map(f, t[p + m]), update_from(p);
    }
    void apply(int l, int r, const F &f) {
        assert(0 <= l && l <= r && r <= n);
        if (l == r) return;
        int li = __builtin_ctz(l + m), ri = __builtin_ctz(r + m);
        push_to(l, li), push_to(r - 1, ri);
        for (int l2 = l + m, r2 = r + m; l2 < r2; l2 >>= 1, r2 >>= 1) {
            if (l2 & 1) all_apply(l2++, f);
            if (r2 & 1) all_apply(--r2, f);
        }
        update_from(l, li), update_from(r - 1, ri);
    }

    template<class G> int max_right(int l, G &&g) {
        assert(0 <= l && l <= n);
        assert(g(M::id()));
        if (l == n) return n;
        push_to(l), l += m;
        T sm = M::id();
        do {
            for (; l % 2 == 0; l >>= 1);
            if (!g(M::op(sm, t[l]))) {
                while (l < m) {
                    push(l), l = l << 1;
                    if (g(M::op(sm, t[l]))) sm = M::op(sm, t[l++]);
                }
                return l - m;
            }
            sm = M::op(sm, t[l++]);
        } while ((l & -l) != l);
        return n;
    }
    template<class G> int min_left(int r, G &&g) {
        assert(0 <= r && r <= n);
        assert(g(M::id()));
        push_to(r - 1), r += m;
        T sm = M::id();
        do {
            for (r--; r > 1 && r % 2; r >>= 1);
            if (!g(M::op(t[r], sm))) {
                while (r < m) {
                    push(r), r = r << 1 | 1;
                    if (g(M::op(t[r], sm))) sm = M::op(t[r--], sm);
                }
                return r + 1 - m;
            }
            sm = M::op(t[r], sm);
        } while ((r & -r) != r);
        return 0;
    }

  private:
    // clang-format off
    static int _lg(int n) { int l = 0; while (1 << l < n) l++; return l; }
    // clang-format on
    void update(int p) { t[p] = M::op(t[p << 1], t[p << 1 | 1]); }
    void update_from(int p, int l = 0) {
        p += m;
        for (int i = l + 1; i <= lg; i++) update(p >> i);
    }
    void all_apply(int p, const F &f) {
        bool ok = M::map(f, t[p]);
        assert(p < m || ok);
        if (p < m) {
            lz[p] = M::comp(f, lz[p]), upd[p] = true;
            if (!ok) push(p), update(p);
        }
    }
    void push(int p) {
        if (!upd[p]) return;
        all_apply(p << 1, lz[p]), all_apply(p << 1 | 1, lz[p]);
        lz[p] = M::fid(), upd[p] = false;
    }
    void push_to(int p, int l = 0) {
        p += m;
        for (int i = lg; i >= l + 1; i--) push(p >> i);
    }
};
#line 2 "math/modint.hpp"

// clang-format off
template<uint32_t m> struct modint {
    static_assert(m >= 1, "Modulus must be in the range [1;2^31)");

    using mint = modint;
    static constexpr bool is_simple = true;

    static constexpr uint32_t mod() noexcept { return m; }
    constexpr modint() noexcept = default;
    constexpr modint(int64_t v) noexcept : _v(uint32_t((v %= m) < 0 ? v + m : v)) {}
    constexpr static mint raw(uint32_t v) noexcept { mint x; return x._v = v, x; }
    template<class T> constexpr explicit operator T() const noexcept { return _v; }

    constexpr mint &operator++() noexcept { return _v = ++_v == mod() ? 0 : _v, *this; }
    constexpr mint &operator--() noexcept { --(_v ? _v : _v = mod()); return *this; }
    constexpr mint operator++(int) noexcept { return exchange(*this, ++mint(*this)); }
    constexpr mint operator--(int) noexcept { return exchange(*this, --mint(*this)); }

    constexpr mint &operator+=(mint rhs) noexcept {
        return _v = int(_v += rhs._v - mod()) < 0 ? _v + mod() : _v, *this;
    }
    constexpr mint &operator-=(mint rhs) noexcept {
        return _v = int(_v -= rhs._v) < 0 ? _v + mod() : _v, *this;
    }
    constexpr mint &operator*=(mint rhs) noexcept {
        return _v = uint64_t(_v) * rhs._v % mod(), *this;
    }
    constexpr mint &operator/=(mint rhs) noexcept {
        return *this = *this * rhs.inv();
    }

    constexpr friend mint operator+(mint l, mint r) noexcept { return l += r; }
    constexpr friend mint operator-(mint l, mint r) noexcept { return l -= r; }
    constexpr friend mint operator*(mint l, mint r) noexcept { return l *= r; }
    constexpr friend mint operator/(mint l, mint r) noexcept { return l /= r; }

    constexpr mint operator+() const noexcept { return *this; }
    constexpr mint operator-() const noexcept { return raw(_v ? mod() - _v : 0); }

    constexpr friend bool operator==(mint l, mint r) noexcept { return l._v == r._v; }
    constexpr friend bool operator!=(mint l, mint r) noexcept { return l._v != r._v; }
    constexpr friend bool operator<(mint l, mint r) noexcept { return l._v < r._v; }

    constexpr mint pow(uint64_t n) const noexcept {
        mint b = *this, res = 1;
        while (n) n & 1 ? res *= b : 0, b *= b, n >>= 1;
        return res;
    }

    constexpr mint inv() const noexcept {
        int a = _v, b = mod(), x = 1, y = 0;
        while (b) {
            x = exchange(y, x - a / b * y);
            a = exchange(b, a % b);
        }
        assert(a == 1);
        return x;
    }

    friend istream &operator>>(istream &is, mint &x) {
        int64_t v{};
        return is >> v, x = v, is;
    }
    friend ostream &operator<<(ostream &os, const mint &x) { return os << x._v; }

  private:
    uint32_t _v = 0;
};
using modint107 = modint<1'000'000'007>;
using modint998 = modint<998'244'353>;
// clang-format on
#line 8 "test/yosupo/ds/segment_tree/range_affine_range_sum.test.cpp"

using mint = modint998;
struct M {
    using T = array<mint, 2>;
    using F = array<mint, 2>;
    static T id() { return {0, 0}; }
    static F fid() { return {1, 0}; }
    static T op(T l, T r) { return {l[0] + r[0], l[1] + r[1]}; }
    static F comp(F l, F r) { return {l[0] * r[0], l[0] * r[1] + l[1]}; }
    static bool map(F f, T &x) {
        x[0] *= f[0];
        x[0] += f[1] * x[1];
        return 1;
    }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    LazySegTree<M> st(n, [](int x) { return cin >> x, M::T{x, 1}; });
    while (q--) {
        int t, l, r, b, c;
        cin >> t >> l >> r;
        if (t == 0) {
            cin >> b >> c;
            st.apply(l, r, {b, c});
        } else {
            cout << st(l, r)[0] << "\n";
        }
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	3 MB
g++	max_random_00	AC	530 ms	16 MB
g++	max_random_01	AC	527 ms	16 MB
g++	max_random_02	AC	522 ms	16 MB
g++	random_00	AC	441 ms	16 MB
g++	random_01	AC	441 ms	16 MB
g++	random_02	AC	298 ms	5 MB
g++	small_00	AC	5 ms	4 MB
g++	small_01	AC	4 ms	4 MB
g++	small_02	AC	4 ms	4 MB
g++	small_03	AC	4 ms	4 MB
g++	small_04	AC	4 ms	4 MB
g++	small_05	AC	4 ms	4 MB
g++	small_06	AC	4 ms	4 MB
g++	small_07	AC	4 ms	4 MB
g++	small_08	AC	4 ms	4 MB
g++	small_09	AC	4 ms	3 MB
g++	small_random_00	AC	5 ms	4 MB
g++	small_random_01	AC	5 ms	4 MB
clang++	example_00	AC	5 ms	4 MB
clang++	max_random_00	AC	543 ms	16 MB
clang++	max_random_01	AC	549 ms	16 MB
clang++	max_random_02	AC	547 ms	16 MB
clang++	random_00	AC	452 ms	16 MB
clang++	random_01	AC	454 ms	16 MB
clang++	random_02	AC	317 ms	5 MB
clang++	small_00	AC	5 ms	4 MB
clang++	small_01	AC	4 ms	4 MB
clang++	small_02	AC	4 ms	4 MB
clang++	small_03	AC	4 ms	4 MB
clang++	small_04	AC	4 ms	4 MB
clang++	small_05	AC	4 ms	4 MB
clang++	small_06	AC	4 ms	4 MB
clang++	small_07	AC	4 ms	4 MB
clang++	small_08	AC	4 ms	4 MB
clang++	small_09	AC	4 ms	4 MB
clang++	small_random_00	AC	5 ms	4 MB
clang++	small_random_01	AC	5 ms	4 MB

Back to top page