Test/LazySegmentTree.test.cpp

View this file on GitHub
Last update: 2021-06-05 13:56:14+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 "Mathematics/ModInt.hpp"
#include "DataStructure/LazySegmentTree.hpp"

#include <cstdio>
#include <utility>

int main() {
	int N, Q;
	scanf("%d %d", &N, &Q);
	
	using mint = ModInt<998244353>;
	
	std::vector<mint> A(N);
	for (int i = 0; i < N; ++i) {
		int a;
		scanf("%d", &a);
		A[i] = a;
	}
	
	using P = std::pair<mint, mint>; // {first} x + {second}

	
	LazySegmentTree<mint, P> seg(A, 0, {1, 0},
		[](const auto & x, const auto & y) { return x + y; },
		[](const auto & x, const auto & e) { return e.first * x + e.second; },
		[](const auto & e1, const auto & e2) { return std::make_pair(e1.first * e2.first, e1.second * e2.first + e2.second); },
		[](const auto & e, auto k) { return std::make_pair(e.first, e.second * k); });
	
	while (Q--) {
		int q, l, r;
		scanf("%d %d %d", &q, &l, &r);
		if (q == 0) {
			int b, c;
			scanf("%d %d", &b, &c);
			seg.update(l, r, {b, c});
		}
		else printf("%d\n", seg.fold(l, r).val());
	}
}

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

#line 1 "Mathematics/ModInt.hpp"



#include <cassert>
#include <iostream>
#include <cstdint>

/**
 * @brief https://tkmst201.github.io/Library/Mathematics/ModInt.hpp
 */
template<int M>
struct ModInt {
	static_assert(M > 0);
	
	using value_type = int;
	using calc_type = std::int_fast64_t;
	
private:
	value_type val_;
	
public:
	constexpr ModInt(calc_type val = 0) : val_(val % M + (val >= 0 ? 0 : M)) {}
	constexpr value_type val() const noexcept { return val_; }
	constexpr static decltype(M) mod() noexcept { return M; }
	
	explicit constexpr operator bool() const noexcept { return val_; }
	constexpr bool operator !() const noexcept { return !static_cast<bool>(*this); }
	constexpr ModInt operator +() const noexcept { return *this; }
	constexpr ModInt operator -() const noexcept { return ModInt(val_ == 0 ? 0 : M - val_); }
	constexpr ModInt operator ++(int) noexcept { ModInt res = *this; ++*this; return res; }
	constexpr ModInt operator --(int) noexcept { ModInt res = *this; --*this; return res; }
	constexpr ModInt & operator ++() noexcept { val_ = val_ + 1 == M ? 0 : val_ + 1; return *this; }
	constexpr ModInt & operator --() noexcept { val_ = val_ == 0 ? M - 1 : val_ - 1; return *this; }
	constexpr ModInt & operator +=(const ModInt & rhs) noexcept { val_ += val_ < M - rhs.val_ ? rhs.val_ : rhs.val_ - M; return *this; }
	constexpr ModInt & operator -=(const ModInt & rhs) noexcept { val_ += val_ >= rhs.val_ ? -rhs.val_ : M - rhs.val_; return *this; }
	constexpr ModInt & operator *=(const ModInt & rhs) noexcept { val_ = static_cast<calc_type>(val_) * rhs.val_ % M; return *this; }
	constexpr ModInt & operator /=(const ModInt & rhs) noexcept { return *this *= rhs.inv(); }
	friend constexpr ModInt operator +(const ModInt & lhs, const ModInt & rhs) noexcept { return ModInt(lhs) += rhs; }
	friend constexpr ModInt operator -(const ModInt & lhs, const ModInt & rhs) noexcept { return ModInt(lhs) -= rhs; }
	friend constexpr ModInt operator *(const ModInt & lhs, const ModInt & rhs) noexcept { return ModInt(lhs) *= rhs; }
	friend constexpr ModInt operator /(const ModInt & lhs, const ModInt & rhs) noexcept { return ModInt(lhs) /= rhs; }
	friend constexpr bool operator ==(const ModInt & lhs, const ModInt & rhs) noexcept { return lhs.val_ == rhs.val_; }
	friend constexpr bool operator !=(const ModInt & lhs, const ModInt & rhs) noexcept { return !(lhs == rhs); }
	friend std::ostream & operator <<(std::ostream & os, const ModInt & rhs) { return os << rhs.val_; }
	friend std::istream & operator >>(std::istream & is, ModInt & rhs) { calc_type x; is >> x; rhs = ModInt(x); return is; }
	
	constexpr ModInt pow(calc_type n) const noexcept {
		ModInt res = 1, x = val_;
		if (n < 0) { x = x.inv(); n = -n; }
		while (n) { if (n & 1) res *= x; x *= x; n >>= 1; }
		return res;
	}
	
	constexpr ModInt inv() const noexcept {
		value_type a = val_, a1 = 1, b = M, b1 = 0;
		while (b > 0) {
			const value_type q = a / b;
			value_type tmp = a - q * b; a = b; b = tmp;
			tmp = a1 - q * b1; a1 = b1; b1 = tmp;
		}
		assert(a == 1);
		if (a1 < 0) a1 += M;
		return a1;
	}
};


#line 1 "DataStructure/LazySegmentTree.hpp"



#include <vector>
#line 6 "DataStructure/LazySegmentTree.hpp"
#include <functional>

/**
 * @brief https://tkmst201.github.io/Library/DataStructure/LazySegmentTree.hpp
 */
template<typename T, typename E>
struct LazySegmentTree {
	using value_type = T;
	using lazy_type = E;
	using size_type = std::size_t;
	using F = std::function<value_type (const value_type &, const value_type &)>;
	using G = std::function<value_type (const value_type &, const lazy_type &)>;
	using H = std::function<lazy_type (const lazy_type &, const lazy_type &)>;
	using P = std::function<lazy_type (const lazy_type &, size_type)>;
	
private:
	size_type n, n_, n_log;
	value_type id_node;
	lazy_type id_lazy;
	F f;
	G g;
	H h;
	P p;
	std::vector<value_type> node;
	std::vector<lazy_type> lazy;

public:
	LazySegmentTree(size_type n, const value_type & id_node, const lazy_type & id_lazy, const F & f, const G & g, const H & h, const P & p = [](const lazy_type & e, size_type k) { return e; })
			: n(n), id_node(id_node), id_lazy(id_lazy), f(f), g(g), h(h), p(p) {
		n_ = 1;
		n_log = 0;
		while (n_ < n) n_ <<= 1, ++n_log;
		node.assign(2 * n_, id_node);
		lazy.assign(2 * n_, id_lazy);
	}
	
	LazySegmentTree(const std::vector<value_type> & v, const value_type & id_node, const lazy_type & id_lazy, const F & f, const G & g, const H & h, const P & p = [](const lazy_type & a, size_type l) { return a; })
			: LazySegmentTree(v.size(), id_node, id_lazy, f, g, h, p) {
		for (size_type i = 0; i < v.size(); ++i) node[i + n_] = v[i];
		for (size_type i = n_ - 1; i > 0; --i) node[i] = f(node[i << 1], node[i << 1 | 1]);
	}
	
	size_type size() const noexcept {
		return n;
	}
	
	void set(size_type k, const value_type & x) noexcept {
		assert(k < size());
		k += n_;
		thrust(k);
		node[k] = x;
		lazy[k] = id_lazy;
		recalc(k);
	}
	
	value_type get(size_type k) noexcept {
		assert(k < size());
		k += n_;
		thrust(k);
		return reflect(k, 1);
	}
	
	void update(size_type l, size_type r, const lazy_type & x) noexcept {
		assert(l <= r);
		assert(r <= size());
		if (l == r) return;
		l += n_;
		r += n_;
		thrust(l);
		thrust(r - 1);
		for (size_type cl = l, cr = r; cl < cr; cl >>= 1, cr >>= 1) {
			if (cl & 1) lazy[cl] = h(lazy[cl], x), ++cl;
			if (cr & 1) --cr, lazy[cr] = h(lazy[cr], x);
		}
		recalc(l);
		recalc(r - 1);
	}
	
	value_type fold(size_type l, size_type r) noexcept {
		assert(l <= r);
		assert(r <= size());
		if (l == r) return id_node;
		l += n_;
		r += n_;
		thrust(l);
		thrust(r - 1);
		value_type vl = id_node, vr = id_node;
		for (size_type w = 1; l < r; l >>= 1, r >>= 1, w <<= 1) {
			if (l & 1) vl = f(vl, reflect(l, w)), ++l;
			if (r & 1) --r, vr = f(reflect(r, w), vr);
		}
		return f(vl, vr);
	}
	
	value_type fold_all() const {
		return reflect(1, n_);
	}
	
private:
	value_type reflect(size_type k, size_type w) const noexcept {
		return lazy[k] == id_lazy ? node[k] : g(node[k], p(lazy[k], w));
	}
	
	void propagate(size_type k, size_type w) noexcept {
		if (lazy[k] == id_lazy) return;
		if ((k << 1) < node.size()) {
			lazy[k << 1] = h(lazy[k << 1], lazy[k]);
			lazy[k << 1 | 1] = h(lazy[k << 1 | 1], lazy[k]);
		}
		node[k] = reflect(k, w);
		lazy[k] = id_lazy;
	}
	
	void recalc(size_type k) noexcept {
		for (size_type i = k >> 1, cw = 1; i > 0; i >>= 1, cw <<= 1)
			node[i] = f(reflect(i << 1, cw), reflect(i << 1 | 1, cw));
	}
	
	void thrust(size_type k) noexcept {
		for (size_type i = n_log, w = n_; i > 0; --i, w >>= 1) propagate(k >> i, w);
	}
};


#line 5 "Test/LazySegmentTree.test.cpp"

#include <cstdio>
#include <utility>

int main() {
	int N, Q;
	scanf("%d %d", &N, &Q);
	
	using mint = ModInt<998244353>;
	
	std::vector<mint> A(N);
	for (int i = 0; i < N; ++i) {
		int a;
		scanf("%d", &a);
		A[i] = a;
	}
	
	using P = std::pair<mint, mint>; // {first} x + {second}

	
	LazySegmentTree<mint, P> seg(A, 0, {1, 0},
		[](const auto & x, const auto & y) { return x + y; },
		[](const auto & x, const auto & e) { return e.first * x + e.second; },
		[](const auto & e1, const auto & e2) { return std::make_pair(e1.first * e2.first, e1.second * e2.first + e2.second); },
		[](const auto & e, auto k) { return std::make_pair(e.first, e.second * k); });
	
	while (Q--) {
		int q, l, r;
		scanf("%d %d %d", &q, &l, &r);
		if (q == 0) {
			int b, c;
			scanf("%d %d", &b, &c);
			seg.update(l, r, {b, c});
		}
		else printf("%d\n", seg.fold(l, r).val());
	}
}