Link to this code:
https://cses.fi/paste/a326c6ec810f0a12fec97e/#include<iostream>
#include<algorithm>
#include<vector>
#include<cassert>
#define TWOPOW(x) ((1)<<(x))
using namespace std;
long long floor_log_2(long long a) {
assert(a > 0);
int k = 0;
while (a > 1) {
a = a / 2;
k = k + 1;
}
return k;
}
bool subseteq(int a,int b,int c,int d) {
return c <= a and b <= d; // c a b d
}
bool disjoint(int a,int b,int c,int d) {
assert(a <= b);
assert(c <= d);
return b < c or d < a;
}
vector<vector<long long>> construct_segment_tree(vector<long long>& s) {
int lh_s = s.size();
int k = floor_log_2(lh_s);
assert(lh_s == TWOPOW(k) ); // lh_s is a power of two
vector<vector<long long>> f_S = vector<vector<long long>>(k + 1); // |[0,k]|= k + 1
int level_size = lh_s;
f_S[0] = s;
level_size = level_size / 2;
int l = 1;
while (l <= k) {
f_S[l] = vector<long long>(level_size);
int j = 0;
while (j < level_size) {
f_S[l][j] = f_S[l - 1][2 * j] + f_S[l - 1][2 * j + 1];
j = j + 1;
}
level_size = level_size / 2;
l = l + 1;
}
return f_S;
}
vector<vector<long long>> construct_empty_segment_tree(int lh_s) {
int k = floor_log_2(lh_s);
assert(lh_s == TWOPOW(k)); // lh_s is a power of two
vector<vector<long long>> seg = vector<vector<long long>>(k + 1);
int level_size = lh_s;
int l = 0;
while (l <= k) {
seg[l] = vector<long long>(level_size);
level_size = level_size / 2;
l = l + 1;
}
return seg;
}
long long get(vector<vector<long long>>& lazy1_f_S, vector<vector<long long>>& lazy2_f_S, vector<vector<long long>>& f_S, int k, int l, int j, long long a, long long b) {
assert(k == (long long)lazy1_f_S.size() - 1);
assert(k == (long long)lazy2_f_S.size() - 1);
assert(k == (long long)f_S.size() - 1);
assert(0 <= l);
assert(l <= k);
assert(0 <= j);
assert(j <= TWOPOW(k-l)-1);
assert(a <= b);
long long c = TWOPOW(l)*(j); // 2^l \times j
long long d = (TWOPOW(l)*(j+1))-1; // 2^l \times (j+1) - 1
assert(c <= d);
if (lazy1_f_S[l][j] != 0) {
f_S[l][j] = f_S[l][j] + (d - c + 1)*lazy1_f_S[l][j];
if (l > 0) {
lazy1_f_S[l-1][2*j] += lazy1_f_S[l][j];
lazy1_f_S[l-1][2*j+1] += + lazy1_f_S[l][j];
}
lazy1_f_S[l][j] = 0;
}
if (lazy2_f_S[l][j] != 0) {
f_S[l][j] = f_S[l][j] + lazy2_f_S[l][j] * ( (d - c + 1) * c + ( ((d - c + 2)*(d - c + 1)) / 2 ) );
if (l > 0) {
lazy2_f_S[l-1][2*j] += lazy2_f_S[l][j];
lazy2_f_S[l-1][2*j+1] += lazy2_f_S[l][j];
}
lazy2_f_S[l][j] = 0;
}
if (disjoint(c,d,a,b)) {
return 0;
}
else {
if (subseteq(c,d,a,b)) {
return f_S[l][j];
}
else {
long long left = get(lazy1_f_S, lazy2_f_S, f_S, k, l-1,2*j,a,b);
long long right = get(lazy1_f_S, lazy2_f_S, f_S, k, l-1,2*j+1,a,b);
return left + right;
}
}
}
void update(vector<vector<long long>>& lazy1_f_S, vector<vector<long long>>& lazy2_f_S, vector<vector<long long>>& f_S, int k, int l, int j, long long a, long long b) {
assert(k == (long long)lazy1_f_S.size() - 1);
assert(k == (long long)lazy2_f_S.size() - 1);
assert(k == (long long)f_S.size() - 1);
assert(0 <= l);
assert(l <= k);
assert(0 <= j);
assert(j <= TWOPOW(k-l)-1);
assert(a <= b);
long long c = TWOPOW(l)*(j); // 2^l \times j
long long d = (TWOPOW(l)*(j+1))-1; // 2^l \times (j+1) - 1
assert(c <= d);
if (not disjoint(c,d,a,b)) {
if (subseteq(c,d,a,b)) {
lazy1_f_S[l][j] = lazy1_f_S[l][j] - a;
lazy2_f_S[l][j] = lazy2_f_S[l][j] + 1;
}
else {
long long s = std::max(c,a);
long long t = std::min(d,b);
assert(s <= t);
f_S[l][j] += (t - s + 1)*(s - a) + (((t - s + 2)*(t - s + 1))/ 2);
update(lazy1_f_S, lazy2_f_S, f_S, k, l-1,2*j,a,b);
update(lazy1_f_S, lazy2_f_S, f_S, k, l-1,2*j+1,a,b);
}
}
}
int main() {
int n;
std::cin >> n;
int q;
std::cin >> q;
int k = floor_log_2(n) + 1;
vector<long long> s(TWOPOW(k)); // 2^{\lfoor log_{2}(n)\rfloor + 1}, i.e. the smallest power of two strictly greater than n
int i = 0;
while (i < n) {
std::cin >> s[i];
i = i + 1;
}
vector<vector<long long>> lazy1_f_S = construct_empty_segment_tree(TWOPOW(k));
vector<vector<long long>> lazy2_f_S = construct_empty_segment_tree(TWOPOW(k));
vector<vector<long long>> f_S = construct_segment_tree(s);
i = 0;
while (i < q) {
int query_type;
long long a;
long long b;
std::cin >> query_type;
std::cin >> a;
std::cin >> b;
if (query_type == 1) {
update(lazy1_f_S, lazy2_f_S, f_S, k, k,0,a-1,b-1);
}
else {
long long val = get(lazy1_f_S, lazy2_f_S, f_S, k, k,0,a-1,b-1);
std::cout << val << std::endl;
}
i = i + 1;
}
}