Range Add Query (RAQ)を図でわかりやすく解説

18250 ワード
C++ 競技プログラミング データ構造 C++ テキストリンク

問題概要

配列Aに対して次の二つのクエリを処理してください

  • $add(s,t,x)$: $A_s,A_{s+1},...,A_t$に$x$を加算する。
  • $get(i)$: $A_i$の値を出力する。

※最初にAの要素は0で初期化します．

入力

n q
query

queryはaddの場合0 s t x，getの場合1 iと入力されます．

制約

n < 100000
q < 100000
x < 1000
1 <= s <= t <= n

大体このような感じです．例えば入力が以下のような場合，下図のように処理します．

入力例

6 5
0 2 4 1
0 3 6 3
1 4
0 4 5 2
1 4

出力例

4
6

赤枠が出力するものです．

解説

例えば配列の2番目から4番目を1足すという処理は1を３回代入するのではなく，配列の2箇所に1と-1を代入します．(配列のサイズが１つ多い事に関しては後述します．)

このようにした後，左から累積和をとると嬉しいことがわかります．(累積和は左から順番に足していくアルゴリズムですが，詳しくはこの記事を参照してください．)

この累積和をとる処理が定数時間でできれば解けそうです．
ここで，セグメント木を使います．セグメント木の詳細な解説に関しては割愛します．（私の簡潔な説明ならこちらの記事や，北海道大学さんのこちらの記事を参照してください．）

セグメント木のできることとして，以下のものがあります．

  • サイズ$N$の配列の区間[$s$, $t$)の和を$O(\log{N})$で計算する

ここで，$s=1$とすると左から$t$個の累積和が取得できます．よってこの問題が解けそうです．そこで先ほどの

配列の2箇所に1と-1を代入します．

をもっと一般化します．

  1. 区間[$s$,$t$)に$x$を足したい時 : $A_s=A_s+x$ と $A_t=A_t-x$ をします．
  2. $A_i$の値が知りたい時 : セグメント木で区間[1, $i$)の和を求めます．

これを先ほどの入力例でやってみます．

左の図の赤枠が右の図の赤枠の合計と一致しています．ここでわかるのですが，$A_t=A_t-x$の処理をする際，$t$は$N$より１つ大きくなることがあります．よって配列のサイズは$N+1$以上にすべきです．

結論

  • $add(s,t,x)$: $A_s,A_{s+1},...,A_t$に$x$を加算する。

が来た時は以下の処理を実行します．(※このクエリの右端は閉区間です．)

A[s] += x;
A[t + 1] -= x;
  • $get(i)$: $A_i$の値を出力する。

このクエリが来た時はセグメント木の機能を使います．

SegmentTree st;
printf("%d\n", st.query(0, i));

実装

コードはこのようになります．セグメント木のコードは長いので隠します．

セグメント木のコード
template <typename T>
class Sum {
public:
    // 単位元
    T unit;

    Sum(void) {
        // 単位元
        unit = 0;
    }

    // 演算関数
    T calc(T a, T b) {
        return a + b; 
    }
};
template <typename T, class MONOID>
class SegmentTree {
public:
    // セグメント木の葉の要素数
    int n;

    // セグメント木
    vector<T> tree;

    // モノイド
    MONOID mono;

    SegmentTree(vector<T> &v) : n(1 << (int)ceil(log2(v.size()))), tree(vector<T>(n << 1)) {
        for(int i = 0; i < v.size(); ++i) {
            update(i, v[i]);
        }
        for(int i = v.size(); i < n; ++i) {
            update(i, mono.unit);
        }
    }

    // k番目の値(0-indexed)をxに変更
    void update(int k, T x) {
        k += n;
        tree[k] = x;
        for(k = k >> 1; k > 0; k >>= 1){
            tree[k] = mono.calc(tree[k << 1 | 0], tree[k << 1 | 1]);
        }
    }

    // [s, t)の最小値(0-indexed)を求める．
    T query(int s, int t) {
        T res = mono.unit;
        s += n;
        t += n;
        while(s < t) {
            if(s & 1) {
                res = mono.calc(res, tree[s++]);
            }
            if(r & 1) {
                res = mono.calc(res, tree[--t]);
            }
            s >>= 1;
            t >>= 1;
        }
        return res;
    }

    T operator[](int k) {
        // st[i]で添字iの要素の値を返す
        if(k - n >= 0 || k < 0) {
            return -INF;
        }
        return tree[tree.size() - n + k];
    }
};


int main() {

    int n, q;
    int query, s, t, x, i;
    cin >> n >> q;

    // セグメント木の作成
    vi v(n + 1);
    SegmentTree<int, Sum<int> > st(v);

    while(q--) {
        scanf("%d", &query);
        if(query) {
            // クエリget(i)の時
            scanf("%d", &i);
            printf("%d\n", st.query(0, i));
        } else {
            // クエリadd(s, t, x)の時
            scanf("%d %d %d", &s, &t, &x);
            // 私のセグメント木の更新処理は加算ではなく代入なので，
            // これはs - 1番目の値をx足しているのと同義です．
            st.update(s - 1, st[s - 1] + x);
            st.update(t, st[t] - x);
        }
    }
}

例題

ほとんど一緒です．最後までご覧いただきありがとうございました．
AOJ-データの集合とクエリ処理2_E-Range Add Query