test/yosupo/range_affine_range_sum.test.cpp

View this file on GitHub
Last update: 2024-03-20 21:25:02+09:00
Problem: https://judge.yosupo.jp/problem/range_affine_range_sum

Depends on

Code

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

#include <bits/stdc++.h>

#include "../../data-structure/segtree/lazy_segment_tree.hpp"
#include "../../math/modint.hpp"
using namespace std;

using mint = Modint<998244353>;

struct M {
    using T = pair<mint, mint>;
    static T id() { return {0, 0}; }
    static T op(T a, T b) { return {a.first + b.first, a.second + b.second}; }
};

struct O {
    using T = pair<mint, mint>;
    static T id() { return {1, 0}; }
    static T op(T a, T b) {
        return {a.first * b.first, a.second * b.first + b.second};
    }
};

M::T act(M::T a, O::T b) {
    return {a.first * b.first + a.second * b.second, a.second};
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N, Q;
    cin >> N >> Q;
    vector<pair<mint, mint>> a(N, {0, 1});
    for (int i = 0; i < N; i++) cin >> a[i].first;
    LazySegmentTree<M, O, act> st(a);
    for (int i = 0; i < Q; i++) {
        int t;
        cin >> t;
        if (t == 0) {
            int l, r, b, c;
            cin >> l >> r >> b >> c;
            st.update(l, r, {b, c});
        } else {
            int l, r;
            cin >> l >> r;
            cout << st.fold(l, r).first << "\n";
        }
    }
}

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

#include <bits/stdc++.h>

#line 3 "data-structure/segtree/lazy_segment_tree.hpp"
#include <bit>
#line 6 "data-structure/segtree/lazy_segment_tree.hpp"

template <typename M, typename O,
          typename M::T (*act)(typename M::T, typename O::T)>
class LazySegmentTree {
    using T = M::T;
    using E = O::T;

   public:
    LazySegmentTree() = default;
    explicit LazySegmentTree(int n)
        : LazySegmentTree(std::vector<T>(n, M::id())) {}
    explicit LazySegmentTree(const std::vector<T>& v)
        : size(std::bit_ceil(v.size())),
          node(2 * size, M::id()),
          lazy(2 * size, O::id()) {
        std::ranges::copy(v, node.begin() + size);
        for (int i = size - 1; i > 0; --i) {
            node[i] = M::op(node[2 * i], node[2 * i + 1]);
        }
    }

    T operator[](int k) { return fold(k, k + 1); }

    void update(int l, int r, const E& x) { update(l, r, x, 1, 0, size); }

    T fold(int l, int r) { return fold(l, r, 1, 0, size); }

    T fold_all() {
        push(1);
        return node[1];
    }

    template <typename F>
    int find_first(int l, F cond) {
        T v = M::id();
        return find_first(l, size, 1, 0, size, v, cond);
    }

    template <typename F>
    int find_last(int r, F cond) {
        T v = M::id();
        return find_last(0, r, 1, 0, size, v, cond);
    }

   private:
    int size;
    std::vector<T> node;
    std::vector<E> lazy;

    void push(int k) {
        if (lazy[k] == O::id()) return;
        if (k < size) {
            lazy[2 * k] = O::op(lazy[2 * k], lazy[k]);
            lazy[2 * k + 1] = O::op(lazy[2 * k + 1], lazy[k]);
        }
        node[k] = act(node[k], lazy[k]);
        lazy[k] = O::id();
    }

    void update(int a, int b, const E& x, int k, int l, int r) {
        push(k);
        if (r <= a || b <= l) return;
        if (a <= l && r <= b) {
            lazy[k] = O::op(lazy[k], x);
            push(k);
            return;
        }
        int m = std::midpoint(l, r);
        update(a, b, x, 2 * k, l, m);
        update(a, b, x, 2 * k + 1, m, r);
        node[k] = M::op(node[2 * k], node[2 * k + 1]);
    }

    T fold(int a, int b, int k, int l, int r) {
        push(k);
        if (r <= a || b <= l) return M::id();
        if (a <= l && r <= b) return node[k];
        int m = std::midpoint(l, r);
        return M::op(fold(a, b, 2 * k, l, m), fold(a, b, 2 * k + 1, m, r));
    }

    template <typename F>
    int find_first(int a, int b, int k, int l, int r, T& v, F cond) {
        push(k);
        if (r <= a) return -1;
        if (b <= l) return l;
        if (a <= l && r <= b && !cond(M::op(v, node[k]))) {
            v = M::op(v, node[k]);
            return -1;
        }
        if (r - l == 1) return r;
        int m = std::midpoint(l, r);
        int res = find_first(a, b, 2 * k, l, m, v, cond);
        if (res != -1) return res;
        return find_first(a, b, 2 * k + 1, m, r, v, cond);
    }

    template <typename F>
    int find_last(int a, int b, int k, int l, int r, T& v, F cond) {
        push(k);
        if (b <= l) return -1;
        if (r <= a) return r;
        if (a <= l && r <= b && !cond(M::op(node[k], v))) {
            v = M::op(node[k], v);
            return -1;
        }
        if (r - l == 1) return l;
        int m = std::midpoint(l, r);
        int res = find_last(a, b, 2 * k + 1, m, r, v, cond);
        if (res != -1) return res;
        return find_last(a, b, 2 * k, l, m, v, cond);
    }
};
#line 4 "math/modint.hpp"

/**
 * @brief Mod int
 */
template <int m>
class Modint {
    using mint = Modint;
    static_assert(m > 0, "Modulus must be positive");

   public:
    static constexpr int mod() { return m; }

    constexpr Modint(long long y = 0) : x(y >= 0 ? y % m : (y % m + m) % m) {}

    constexpr int val() const { return x; }

    constexpr mint& operator+=(const mint& r) {
        if ((x += r.x) >= m) x -= m;
        return *this;
    }
    constexpr mint& operator-=(const mint& r) {
        if ((x += m - r.x) >= m) x -= m;
        return *this;
    }
    constexpr mint& operator*=(const mint& r) {
        x = static_cast<int>(1LL * x * r.x % m);
        return *this;
    }
    constexpr mint& operator/=(const mint& r) { return *this *= r.inv(); }

    constexpr bool operator==(const mint& r) const { return x == r.x; }

    constexpr mint operator+() const { return *this; }
    constexpr mint operator-() const { return mint(-x); }

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

    constexpr mint inv() const {
        int a = x, b = m, u = 1, v = 0;
        while (b > 0) {
            int t = a / b;
            std::swap(a -= t * b, b);
            std::swap(u -= t * v, v);
        }
        return mint(u);
    }

    constexpr mint pow(long long n) const {
        mint ret(1), mul(x);
        while (n > 0) {
            if (n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }

    friend std::ostream& operator<<(std::ostream& os, const mint& r) {
        return os << r.x;
    }

    friend std::istream& operator>>(std::istream& is, mint& r) {
        long long t;
        is >> t;
        r = mint(t);
        return is;
    }

   private:
    int x;
};
#line 7 "test/yosupo/range_affine_range_sum.test.cpp"
using namespace std;

using mint = Modint<998244353>;

struct M {
    using T = pair<mint, mint>;
    static T id() { return {0, 0}; }
    static T op(T a, T b) { return {a.first + b.first, a.second + b.second}; }
};

struct O {
    using T = pair<mint, mint>;
    static T id() { return {1, 0}; }
    static T op(T a, T b) {
        return {a.first * b.first, a.second * b.first + b.second};
    }
};

M::T act(M::T a, O::T b) {
    return {a.first * b.first + a.second * b.second, a.second};
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N, Q;
    cin >> N >> Q;
    vector<pair<mint, mint>> a(N, {0, 1});
    for (int i = 0; i < N; i++) cin >> a[i].first;
    LazySegmentTree<M, O, act> st(a);
    for (int i = 0; i < Q; i++) {
        int t;
        cin >> t;
        if (t == 0) {
            int l, r, b, c;
            cin >> l >> r >> b >> c;
            st.update(l, r, {b, c});
        } else {
            int l, r;
            cin >> l >> r;
            cout << st.fold(l, r).first << "\n";
        }
    }
}