题解 SP11840 【SEGSQRSS - Sum of Squares with Segment Tree】

MVP_Harry · 2020-08-03 12:14:13 · 题解

居然做到了一个没人做过的题

这题感觉还是很不错的，做完之后对线段树的理解又加深了一些，建议添加标签线段树，难度在绿题和蓝题之间（参考线段树模板2）

在这题中，线段树一个节点要保存四个值：sum, square, change, add，分别代表和，平方和和两个lazy tag。不难发现\displaystyle \sum_{i = x}^{y} (a_i + k)^2 = \displaystyle \sum_{i = x}^{y} (a_i^2 + 2a_ik) + (y-x+1)k^2可以用sum维护。

对于lazy tag来说其实自己想一下就出来了，可以对着样例模拟一下～

代码如下：（其实交的时候太激动了没有开long long但还是过了)

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define rep(i, m, n) for(int i = m; i <= n; i++)
#define per(i, m, n) for(int i = m; i >= n; i--)
#define pb push_back
#define DEBUG(x) cout<<x<<endl;
const int INF = 0x3f3f3f3f;

const int maxn = 1e5 + 5;

struct SegmentTree {
    int sum, square, change, add;
} T[maxn << 2];

int N, Q, TIME, A[maxn];

inline void build(int s, int t, int p) { //建树
    if (s == t) {
        T[p].sum = A[s], T[p].square = A[s] * A[s];
        return ;
    } 
    int m = (s + t) >> 1;
    build(s, m, p << 1), build(m + 1, t, p << 1 | 1);
    T[p].sum = T[p << 1].sum + T[p << 1 | 1].sum;
    T[p].square = T[p << 1].square + T[p << 1 | 1].square;
}

inline void pushdown(int s, int t, int p) { //维护lazy tag 
    int m = (s + t) >> 1;
    if (T[p].change) {
        T[p << 1].change = T[p].change, T[p << 1 | 1].change = T[p].change;
        T[p << 1].add = T[p].add, T[p << 1 | 1].add = T[p].add;
        T[p << 1].sum = (m - s + 1) * (T[p].change + T[p].add), T[p << 1 | 1].sum = (t - m) * (T[p].change + T[p].add);
        T[p << 1].square = (m - s + 1) * (T[p].change + T[p].add) * (T[p].change + T[p].add), T[p << 1 | 1].square = (t - m) * (T[p].change + T[p].add) * (T[p].change + T[p].add);
    } else {
        T[p << 1].add += T[p].add, T[p << 1 | 1].add += T[p].add;
        T[p << 1].square += ((m - s + 1) * T[p].add * T[p].add + 2 * T[p].add * T[p << 1].sum), T[p << 1 | 1].square += ((t - m) * T[p].add * T[p].add + 2 * T[p].add * T[p << 1 | 1].sum);
        T[p << 1].sum += (m - s + 1) * T[p].add, T[p << 1 | 1].sum += (t - m) * T[p].add;
    }
}

inline void update1(int l, int r, int x, int s, int t, int p) { //increment values in [l, r] by x
    if (l <= s && r >= t) {
        T[p].square += ((s - t + 1) * x * x + 2 * x * T[p].sum), T[p].sum += (s - t + 1) * x, T[p].add += x;
        return ;
    }
    int m = (s + t) >> 1;
    pushdown(s, t, p);
    if (l <= m) update1(l, r, x, s, m, p << 1);
    if (r > m) update1(l, r, x, m + 1, t, p << 1 | 1);
    T[p].square = T[p << 1].square + T[p << 1 | 1].square;
    T[p].sum = T[p << 1].sum + T[p << 1 | 1].sum;
}

inline void update2(int l, int r, int x, int s, int t, int p) { //更改为x
    if (l <= s && r >= t) {
        T[p].square = (t - s + 1) * x * x, T[p].sum = (t - s + 1) * x, T[p].add = 0, T[p].change = x;
        return ;
    }
    int m = (s + t) >> 1;
    pushdown(s, t, p);
    if (l <= m) update2(l, r, x, s, m, p << 1);
    if (r > m) update2(l, r, x, m + 1, t, p << 1 | 1);
    T[p].square = T[p << 1].square + T[p << 1 | 1].square;
    T[p].sum = T[p << 1].sum + T[p << 1 | 1].sum;
}

inline int query(int l, int r, int s, int t, int p) { //查询
    if (l <= s && r >= t) 
        return T[p].square;
    int m = (s + t) >> 1, S = 0;
    if (l <= m) S += query(l, r, s, m, p << 1);
    if (r > m) S += query(l, r, m + 1, t, p << 1 | 1);
    return S;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int cnt = 0;
    cin >> TIME;
    while (TIME--) {
        cout << "Case " << ++cnt << ":" << endl; 
        cin >> N >> Q;
        for (int i = 1; i <= N; i++) cin >> A[i];
        build(1, N, 1);
        int x, c1, c2, k;
        while (Q--) {
            cin >> x;
            if (x == 0) {
                cin >> c1 >> c2 >> k;
                update2(c1, c2, k, 1, N, 1);
            } else if (x == 1) {
                cin >> c1 >> c2 >> k;
                update1(c1, c2, k, 1, N, 1);
            } else {
                cin >> c1 >> c2;
                cout << query(c1, c2, 1, N, 1) << endl;
            }
        }
    }
    return 0;
}