AtCoder ABC332 F - Random Update Query: LazySegmentTreeメモ

LazySegtreeで考えるメモです。

前提知識: Lazy Segtreeの気持ち

Lazy Segtreeは各範囲に対応する値(区間をprodしたもの)と各範囲に対する演算を持ちます。mapping(F f, S x)はあるノードの持つ演算$f$をこのノードの持つ値xに作用させ、composition(F f, F g)は親の演算fと子の持つ演算をf+gで合成します。compositionは親から子に情報を伝搬した後、演算の単位元であるid()とします。

問題の考え方

$[l,r]$に含まれる$i$の要素$a_i$を考えた時、区間[l,r]の数を$c = r - l + 1$とするとそのクエリ後の$a_i$の期待値は$ \frac{(c-1)}{c} a_i + \frac{1}{c} x $ となります。$c$と$x$はクエリ毎に定数となるので$A \cdot a_i + B$として良いです。
この更新された要素を範囲に含む新しいクエリ$l', r', x'$が与えられた時を考えます。同じように$c' = r' - l' + 1$とすると、その要素の期待値は$ \frac{(c'-1)}{c'} (A \cdot a_i +B) + \frac{1}{c'} x' $となり、$c'$と$x'$は先ほどと同様に定数となるので$ C (A \cdot a_i +B) + D = A \cdot C \cdot a_i + B \cdot C + D = (A \cdot C ) a_i + (B \cdot C + D) $となります。

つまり、ある要素$a_i$はその要素に適応される$l,r,x$がわかれば最終的な期待値$p_i$は$p_i = A \cdot a_i + B$の形で表せます。

各クエリによって範囲の要素の$A, B$を高速に更新できれば良さそうです。期待値$p_i$は何度クエリで処理されても$p_i = A \cdot a_i + B$の形で表せるということはこの演算に閉じており、写像を合成してもこの形で表せるということです。これがLazy Segtreeにおいてcomposition（lazyを下のlazyに伝える）ことなります。

まず、この演算をノードの値に適応する関数$mapping$を考えます。

struct vtype{ modint998244353 a, b; }; /* 演算を載せる構造体 */
using S = modint998244353; /* val の形 */
using F = vtype; /* lazyの形 */
/* 上位のlazy(f)を1つ子のval(x)に伝搬させる演算。ax + b,valにしたい */
S mapping(F f, S x){ return (S)f.a * x + f.b ; }

Lazyを伝搬する関数$composition$は次の通りとなります。ここで、f,gは可換ではなく、親fを子gに伝えていることに注意してください。

/* 子gに親fを伝える。(f,g) は f(g(x))*/
// f.a, f.b, g.a, g.bがそれぞれC,D,A,Bに対応します。A,B,C,Dに対応しないことを注意してください。
F composition(F f, F g){ return F{f.a * g.a, f.a * g.b + f.b }; }

LazySegtreeでは範囲内のノードにLazyを伝搬して値を更新します。ノード自身の初期値と演算に関する単位元が必要になります。

/* 区間, prod: 使わないのでなんでもいい */
auto op = [](S x, S y){ return x+y; };

/* valの初期値 = 0 */
auto e = []{ return (S)0; };

/* lazyを下にmappingで伝搬させても、「valを変更しないよう」なF
 * ax + bのcomposition(単位元,y)になるので {1,0}とすれば{1 * y.a, 0 * 1 + y.b } = y */
F id(){ return (F){1, 0}; }

コメントの通りですが、Fの単位元とは伝搬して子のvalを変更しないようなFであり、子のlazyを変更しないではないことに注意します。

実装

struct vtype{ /* 乗せるタイプ */
    modint998244353 a;
    modint998244353 b;
};
using S = modint998244353; /* val の形 */
using F = vtype; /* lazyの形 */
auto op = [](S x, S y){ return x+y; };
auto e = []{ return (S)0; };
S mapping(F f, S x){ return (S)f.a * x + f.b ; }
F composition(F f, F g){ return F{f.a * g.a, f.a * g.b + f.b }; }
y.a, 0 * 1 + y.b } = y */
F id(){ return (F){1, 0}; }
int main(){
    int n, m; cin >> n >> m;
    std::vector<S> v(n);
    REP(i, n) {
        int x; cin >> x;
        v.at(i) = x;
    }

    atcoder::lazy_segtree<S, op, e, F, mapping, composition, id> seg(v);
    REP(i, m){
        int l, r, x;
        cin >> l >> r >> x;
        l--; r--;
        modint998244353 cnt = r - l + 1;
        seg.apply(l, r+1, F{ (cnt-1) / cnt, (modint998244353)x / cnt });
    }

    REP(i, n){
        cout << seg.get(i).val() << " ";
    }
    cout << endl;

}