yukicoder : No.1099 Range Square Sum

面倒だけど典型的な問題。
https://yukicoder.me/problems/no/1099

問題

配列Aが与えられる。
以下のクエリに順次答えよ。

• 指定された区間内の要素に同じ値kを加算する
• 指定された区間内の要素の二乗和を答える

解法

区間に同じ値を加算すると、二乗和がどうなるか考える。
数列Aに対し、B[i]=A[i]+kで構築したBの二乗和を考えると、

と書ける。
よって、遅延セグメント木を書き、各ノードにこれまで加算した数の総和、区間の和、区間の二乗和を持たせよう。
ノード内の一部に対する加算や二乗和を求めるクエリが来たとき、加算分を下位ノードに伝搬させていく。

template<class V,int NV> class SegTree_1 {
public:
	vector<V> add;
	vector<V> sum1;
	vector<V> sum2;
	
	SegTree_1(){
		add.resize(NV*2);
		sum1.resize(NV*2);
		sum2.resize(NV*2);
	};
	
	V getval(int x,int y,int l,int r,int k) {
		if(r<=x || y<=l) return 0;
		if(x<=l && r<=y) return sum2[k];
		if(add[k]) {
			update(l,(l+r)/2,l,(l+r)/2,k*2,add[k]);
			update((l+r)/2,r,(l+r)/2,r,k*2+1,add[k]);
			add[k]=0;
			sum1[k]=sum1[2*k]+sum1[2*k+1];
			sum2[k]=sum2[2*k]+sum2[2*k+1];
		}
		
		return getval(x,y,l,(l+r)/2,k*2)+getval(x,y,(l+r)/2,r,k*2+1);
	}
	V getval(int x,int y) { return getval(x,y,0,NV,1);}
	
	void update(int x,int y,int l,int r,int k,int v) {
		if(r<=x || y<=l) return;
		if(v==0) return;
		if(x<=l && r<=y) {
			add[k]+=v;
			sum2[k]+=2LL*v*sum1[k]+1LL*(r-l)*v*v;
			sum1[k]+=1LL*v*(r-l);
			return;
		}
		
		update(l,(l+r),l,(l+r)/2,k*2,add[k]);
		update((l+r)/2,r,(l+r)/2,r,k*2+1,add[k]);
		update(x,y,l,(l+r)/2,k*2,v);
		update(x,y,(l+r)/2,r,k*2+1,v);
		add[k]=0;
		sum1[k]=sum1[2*k]+sum1[2*k+1];
		sum2[k]=sum2[2*k]+sum2[2*k+1];
	}
	void update(int x,int y,int v) { update(x,y,0,NV,1,v);}
};
SegTree_1<ll,1<<20> st;

int N;
int Q;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>x;
		st.update(i+1,i+2,x);
	}
	cin>>Q;
	while(Q--) {
		cin>>i>>l>>r;
		if(i==1) {
			cin>>x;
			st.update(l,r+1,x);
		}
		else {
			cout<<st.getval(l,r+1)<<endl;
		}
	}
}

まとめ

これ★3.5ならFは★4でいいんじゃ…。