SEGSQRSS - Sum of Squares with Segment Tree 题解

· · 题解

SEGSQRSS - Sum of Squares with Segment Tree

本题是一道线段树板子，建议评绿。

思路：

我们需要完成区间加，区间修改和求区间平方和的操作。

对于前两个操作，我们可以维护 3 个懒标记，分别表示区间加的值，区间是否有修改，区间修改为的值。当然，在这题中，加数和修改的绝对值都不超过一千，所以可以初始时将修改赋为一个绝对值大于一千的数来表示没有被修改过。

对于第三个操作，我们观察一下在加法后区间的平方和发生了什么变化：（下文中的 c 表示要加上的数或要修改为的数）

加法：

\sum_{i=1}^{n}(a_i+c)^2=\sum_{i=1}^n(a_i^2+2a_ic+c^2)=\sum_{i=1}^na_i^2+\sum_{i=1}^n2a_ic+\sum_{i=1}^nc^2=\sum_{i=1}^na_i^2+2c\sum_{i=1}^na_i+nc^2

对于修改操作则更简单，此时平方和将变为 nc^2。

我们发现，为了 O(1) 得到新的平方和，我们需要在维护平方和的同时维护序列的区间和。

因此，我们需要在线段树中维护懒标记，区间和，区间平方和，再按照上文所说的方法更新就可以了。

做法及技巧：

注意一下下放懒标记的顺序，应该先下放修改标记再下放加法标记。

这是因为更新修改标记会清空加法标记。也就是说如果一个区间同时有加法和修改标记，那么这个区间是修改后再被加的，因此要先下放修改标记，防止就加法标记被清空。

多测，记得清空。

代码：

#include <bits/stdc++.h>
using namespace std;
const int N=100100;
typedef long long ll;

int n,m,T,in1,in2,in3,op,num;
int inp[N];//输入的序列

struct STn{int l,r,len;ll t1,t2,t3,sum,sum2;};
//区间左，右端点，区间长度，三个懒标记，和，平方和
struct ST{
    STn a[N<<2];
    void push_up(int p){//上传
        a[p].sum=a[p<<1].sum+a[p<<1|1].sum;
        a[p].sum2=a[p<<1].sum2+a[p<<1|1].sum2;
        return ;
    }
    void change_t(int p,int k){//修改某一区间
        a[p].t1=0;a[p].t2=1;a[p].t3=k;
        a[p].sum2=a[p].len*k*k;a[p].sum=a[p].len*k;
        return ;
    }
    void add_t(int p,int k){
        a[p].t1+=k;
        a[p].sum2=a[p].sum2+2*k*a[p].sum+a[p].len*k*k;
        a[p].sum+=k*a[p].len;//数据貌似很水，我第一发忘加这句话居然也过了
        return ;
    }
    void push_down(int p){
        if(a[p].t2){//先下放修改
            change_t(p<<1,a[p].t3);
            change_t(p<<1|1,a[p].t3);
            a[p].t2=a[p].t3=0;
        }
        if(a[p].t1){
            add_t(p<<1,a[p].t1);
            add_t(p<<1|1,a[p].t1);
            a[p].t1=0;
        }
        return ;
    }
    void build(int p,int l,int r){
        a[p].l=l;a[p].r=r;a[p].len=r-l+1;
        a[p].t1=a[p].t2=a[p].t3=0;//清空标记
        if(a[p].l==a[p].r){
            a[p].sum=inp[a[p].l];
            a[p].sum2=a[p].sum*a[p].sum;
            return ;
        }
        int mid=(a[p].l+a[p].r)>>1;
        build(p<<1,l,mid);
        build(p<<1|1,mid+1,r);
        push_up(p);return ;
    }
    void add_and_change(int p,int l,int r,int k,int f){//整合成了一个函数
        if(l<=a[p].l&&a[p].r<=r){
            if(f==1) add_t(p,k);
            if(f==2) change_t(p,k);
            return ;
        }
        push_down(p);int mid=(a[p].l+a[p].r)>>1;
        if(l<=mid) add_and_change(p<<1,l,r,k,f);
        if(r>mid) add_and_change(p<<1|1,l,r,k,f);
        push_up(p);return ;
    }
    ll ask(int p,int l,int r){
        if(l<=a[p].l&&a[p].r<=r) return a[p].sum2;
        push_down(p);int mid=(a[p].l+a[p].r)>>1;
        if(r<=mid) return ask(p<<1,l,r);
        if(l>mid) return ask(p<<1|1,l,r);
        return ask(p<<1,l,r)+ask(p<<1|1,l,r);
    }
}tree;

int main(){
    scanf("%d",&T);
    while(T--){
        printf("Case %d:\n",++num);
        scanf("%d%d",&n,&m);
        for(int i=1;i<=n;i++) 
            scanf("%d",&inp[i]);
        tree.build(1,1,n);
        while(m--){
            scanf("%d%d%d",&op,&in1,&in2);
            if(op==0){
                scanf("%d",&in3);
                tree.add_and_change(1,in1,in2,in3,2);//修改
            }   
            if(op==1){
                scanf("%d",&in3);
                tree.add_and_change(1,in1,in2,in3,1);//区间加
            }
            if(op==2)
                printf("%lld\n",tree.ask(1,in1,in2));
        }
    }
    return 0;
}

其他：

与该题类似的题目：P1471 方差。