(セグメント木・区間加算区間和・遅延評価)AOJ-DSL_2-G-Range Query – RSQ and RAQ

問題URL→http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_G&lang=jp

このブログを見て遅延評価セグメント木を実装した。

タイトル通り、この実装では区間加算と区間和のクエリに対応している。

//*
#include <bits/stdc++.h>
//*/
/*
#include <iostream>
#include <vector>
#include <string.h>
//*/
#define FOR(i,bg,ed) for(ll i=(bg);i<(ed);i++)
#define REP(i,n) FOR(i,0,n)
#define MOD 1000000007

#define int long long
#ifdef int
    const long long INF = LLONG_MAX / 10;
#else
    const int INF = 1010101010;
#endif

using namespace std;
typedef long long ll;
typedef vector<vector<ll>> mat;
typedef pair<int, int> P;
//typedef pair<double, double> P;

//遅延評価セグメント木(区間加算区間和)
struct LazySegmentTree
{
private:
    int n;
    vector<ll> node, lazy;

public:
    //初期化
    LazySegmentTree(vector<ll> v)
    {
        int sz = (int)v.size();
        n = 1;
        while (n < sz) n *= 2;
        node.resize(2 * n - 1);
        lazy.resize(2 * n - 1, 0);

        for (int i=0; i<sz; i++) node[i+(n-1)] = v[i];
        for (int i=n-2; i>=0; i--) node[i] = node[i*2+1] + node[i*2+2];
    }

    //k番目のノードについて遅延評価を行う
    void eval(int k, int l, int r)
    {
        //遅延配列が空でない場合、自ノード及び子ノードへの
        //値の伝播が起こる
        if (lazy[k] != 0) {
            node[k] += lazy[k];

            //最下段かどうかのチェックをする
            //子ノードは親ノードの1/2の範囲であるため、
            //伝播させるときは半分にする
            if (r - l > 1) {
                lazy[2*k+1] += lazy[k] / 2;
                lazy[2*k+2] += lazy[k] / 2;
            }

            //伝播が終わったので、自ノードの遅延配列を空にする
            lazy[k] = 0;
        }
    }

    //区間加算
    void add(int a, int b, ll x, int k=0, int l=0, int r=-1)
    {
        if (r < 0) r = n;

        //k番目のノードに対して遅延評価を行う
        eval(k, l, r);

        //範囲外なら何もしない
        if (b <= l || r <= a) return;

        //完全に被覆しているならば、遅延配列に値を入れた後に評価
        if (a <= l && r <= b) {
            lazy[k] += (r - l) * x;
            eval(k, l, r);
        }

        //そうでないならば、子ノードの値を再帰的に計算して、
        //計算済みの値をもらってくる
        else {
            add(a, b, x, 2 * k + 1, l, (l + r) / 2);
            add(a, b, x, 2 * k + 2, (l + r) / 2, r);
            node[k] = node[2*k+1] + node[2*k+2];
        }
    }

    //区間和の取得
    ll getsum(int a, int b, int k=0, int l=0, int r=-1)
    {
        if (r < 0) r = n;

        //関数が呼び出されたらまず評価
        eval(k, l, r);

        if (b <= l || r <= a) return 0;
        if (a <= l && r <= b) return node[k];
        ll vl = getsum(a, b, 2*k+1, l, (l+r)/2);
        ll vr = getsum(a, b, 2*k+2, (l+r)/2, r);
        return vl + vr;
    }
};

signed main()
{
    int n, q;

    cin >> n >> q;
    vector<ll> a(n, 0);
    LazySegmentTree lst(a);

    REP(i,q) {
        int type;
        cin >> type;
        if (type == 0) {
            int s, t, x;
            cin >> s >> t >> x;
            lst.add(s - 1, t, x);
        } else {
            int s, t;
            cin >> s >> t;
            cout << lst.getsum(s - 1, t) << endl;
        }
    }
}