🐈

ある範囲内の要素の値を切り替える区間クエリと区間和を実現する遅延セグメント木(AtCoder Library使用)

2022/02/05に公開

ある範囲内の要素の値を一方から他方へ変更する区間変更クエリと全体の総和をとるクエリを実現する遅延セグメント木を、AtCoder LibraryのLazySegmentTreeを用いて実装する方法について。

AtCoder LibraryのLazySegmentTree

https://atcoder.github.io/ac-library/document_ja/lazysegtree.html

どういう問題が解けるか

次のような問題が解ける。

問題

数列A = \{a_0, a_1, ..., a_{N-1}\}が与えられる。これと同じ数列、つまり0\leq i < Nの時a_i = b_iとなる数列B = \{b_0, b_1, ..., b_{N-1}\}を用意する。以下の二つのクエリを処理してください。

  1. i = l, l+1, ..., r-1に対して、a_i = b_iならばa_i\leftarrow 0とし、a_i\neq b_iならばa_i\leftarrow b_iと変更する。
  2. i = l, l+1, ..., r-1に対して、\sum_{i=l}^{r-1}a_iを出力する。

制約: 1\leq N\leq2×10^5, |a_i| \leq 10^9

解法

遅延セグメント木を用いる。
クエリ1の区間変更に関しては、それぞれの値の符号を反転させるようにする。
クエリ2の区間和は、セグメント木上の区間和をとってきた後に、元の配列の値の該当する区間の総和を足し、2で割ることで求められる。
元の配列の値の総和は累積和を用いてO(1)で求められる。

クエリの回数をQとすると、各クエリはO(log N)で処理できるので、全体の計算量はO(QlogN)となる。

具体例

具体例は以下。seg[i]はセグメント木上の値、A[i]は元の配列の値を表す。

区間クエリに対して求める値は、列和の総和(もしくは行和の総和)である右下の赤いマスの値の半分である。各列の和を見ると、元の配列の値との和を取ることで、反転している要素の値はゼロになり、そのままの値は元の値の2倍になっていることがわかる。なのでこの列和の総和を2で割れば、反転していない要素の総和が求まる。

初期値

[0, 3) に対して反転を行う。

[2, 5) に対して反転を行う。

実装例

#include <bits/stdc++.h>
#include <atcoder/lazysegtree>
using namespace std;

namespace RangeToggleSumLazySegtree {
using S = long long;
using F = bool;
S op(S a, S b) { return a + b; }
S e() { return S(0); }
S mapping(F f, S x) { return f ? x : -x; }
F composition(F f, F g) { return f == g; }
F id() { return true; }

using _segtree = atcoder::lazy_segtree<S, op, e, F, mapping, composition, id>;

struct SegTree {
   private:
    vector<S> v, cum_sum;
    _segtree seg;

   public:
    SegTree(vector<S> v) : v(v) {
        seg = _segtree(v);

        cum_sum = {0};
        for (S &e : v) cum_sum.push_back(cum_sum.back() + e);
    }

    S prod(int l, int r) {
        return (cum_sum[r] - cum_sum[l] + seg.prod(l, r)) / 2;
    }

    S all_prod() { return prod(0, v.size()); }

    void apply(int l, int r) {
        seg.apply(l, r, false);
    }
};
}  // namespace RangeToggleSumLazySegtree

使用例

int main() {
    vector<RangeToggleSumLazySegtree::S> a = {3, 2, 5, 8, 4, 1};
    const int n = a.size();

    RangeToggleSumLazySegtree::SegTree seg(a);

    // 最初は全てその値
    cout << seg.all_prod() << endl;  // 23

    // [0, n) の値を切り替える
    seg.apply(0, n);  // -> a = {0, 0, 0, 0, 0, 0}

    cout << seg.prod(0, n) << endl;  // 0

    // [0, 3) の値を切り替える
    seg.apply(0, 3);  // -> a = {3, 2, 5, 0, 0, 0}

    cout << seg.all_prod() << endl;  // 10
    cout << seg.prod(2, n) << endl;  // 5

    // [1, 5) の値を切り替える
    seg.apply(1, 5);  // -> a = {3, 0, 0, 8, 4, 0}

    cout << seg.all_prod() << endl;  // 15
    cout << seg.prod(3, 5) << endl;  // 12

    return 0;
}

提出

PAST 第7回 N - モノクロデザイン/Monochrome Design に対する提出。

https://atcoder.jp/contests/past202107-open/submissions/29053019

もっといいやり方、ありそう?

あれば教えてください。

Discussion