tests/lazy_segment_tree.test.cpp

View this file on GitHub
Last update: 2024-02-01 17:09:56-05: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 "../content/data_structures/lazy_segment_tree.h"

using pll = pair<ll, ll>;
constexpr ll mod = 998244353;

int main(){
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cin.exceptions(cin.failbit);
	
	int n, q;
	cin>>n>>q;
	
	lazy_segment_tree<ll, pll> ST;
	ST.init(n+5, 0, pll(1, 0),
		[](auto x, auto y){ return (x+y) % mod; },
		[](auto x, auto y, int l, int r){ return (x*y.first + (r-l+1)*y.second) % mod; },
		[](auto x, auto y){ return pll(x.first*y.first % mod, (y.first*x.second + y.second) % mod); }
		);
	
	for(int i = 0; i < n; i++){
		ll a;
		cin>>a;
		ST.update(i, i, {0, a});
	}
	while(q--){
		int op;
		cin>>op;
		if(op == 0){
			int l, r; ll b, c;
			cin>>l>>r>>b>>c;
			ST.update(l, r-1, {b, c});
		}
		else{
			int l, r;
			cin>>l>>r;
			auto a = ST.query(l, r-1);
			cout<<a<<'\n';
		}
	}
}

#line 1 "tests/lazy_segment_tree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/range_affine_range_sum"
#line 1 "content/utils/template.h"
/**
 * @brief My starter code
 */

#include <bits/stdc++.h>
#define all(x) (x).begin(), (x).end()
#define makeunique(x) (x).erase(unique((x).begin(), (x).end()), (x).end());

#ifdef LOCAL
template<typename T> void pr(T a){std::cerr<<a<<std::endl;}
template<typename T, typename... Args> void pr(T a, Args... args){std::cerr<<a<<' ',pr(args...);}
#else
template<typename... Args> void pr(Args... args){}
#endif

using namespace std;
using ll = long long;
using pii = pair<int, int>;
#line 3 "content/data_structures/lazy_segment_tree.h"

/**
 * @brief Lazy Segment Tree
 *
 * todo: make "find first in range" (from cp algorthms) function
 * to do that, I need a function to check if f(x, current) is good enough, and I also need combine(current, val[rt])
 *
 */

template<class T, class L> struct lazy_segment_tree{
	int SZ;
	vector<T> val;
	vector<L> lp;
	function<T(T, T)> merge;
	function<T(T, L, int, int)> apply;
	function<L(L, L)> merge_lazy;
	
	void init(int n, T def, L ldef, function<T(T, T)> f, function<T(T, L, int, int)> g, function<L(L, L)> h){
		SZ = n;
		while(n&(n-1)) n++; // next largest power of 2
		val.resize(2*n, def);
		lp.resize(size(val), ldef);
		merge = f;
		apply = g;
		merge_lazy = h;
	}

	#define lc (rt<<1)
	#define rc (rt<<1|1)
	#define nm ((nl+nr)/2)

	void push(int nl, int nr, int rt){
		L &x = lp[rt];
		val[lc] = apply(val[lc], x, nl, nm);
		lp[lc] = merge_lazy(lp[lc], x);
		val[rc] = apply(val[rc], x, nm+1, nr);
		lp[rc] = merge_lazy(lp[rc], x);
		x = lp[0];
	}
	
	T _query(int l, int r, int nl, int nr, int rt){
		if(r < nl or nr < l) return val[0];
		if(l <= nl and nr <= r) return val[rt];
		push(nl, nr, rt);
		return merge(_query(l, r, nl, nm, lc), _query(l, r, nm+1, nr, rc));
	}
	T query(int l, int r){ return _query(l, r, 0, SZ-1, 1); }
	
	void _update(int l, int r, L x, int nl, int nr, int rt){
		if(r < nl or nr < l) return;
		if(l <= nl and nr <= r){
			val[rt] = apply(val[rt], x, nl, nr);
			lp[rt] = merge_lazy(lp[rt], x);
			return;
		}
		push(nl, nr, rt);
		_update(l, r, x, nl, nm, lc);
		_update(l, r, x, nm+1, nr, rc);
		val[rt] = merge(val[lc], val[rc]);
	}
	void update(int l, int r, L x){ _update(l, r, x, 0, SZ-1, 1); };

	#undef lc
	#undef rc
	#undef nm
};
#line 3 "tests/lazy_segment_tree.test.cpp"

using pll = pair<ll, ll>;
constexpr ll mod = 998244353;

int main(){
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cin.exceptions(cin.failbit);
	
	int n, q;
	cin>>n>>q;
	
	lazy_segment_tree<ll, pll> ST;
	ST.init(n+5, 0, pll(1, 0),
		[](auto x, auto y){ return (x+y) % mod; },
		[](auto x, auto y, int l, int r){ return (x*y.first + (r-l+1)*y.second) % mod; },
		[](auto x, auto y){ return pll(x.first*y.first % mod, (y.first*x.second + y.second) % mod); }
		);
	
	for(int i = 0; i < n; i++){
		ll a;
		cin>>a;
		ST.update(i, i, {0, a});
	}
	while(q--){
		int op;
		cin>>op;
		if(op == 0){
			int l, r; ll b, c;
			cin>>l>>r>>b>>c;
			ST.update(l, r-1, {b, c});
		}
		else{
			int l, r;
			cin>>l>>r;
			auto a = ST.query(l, r-1);
			cout<<a<<'\n';
		}
	}
}