ソラで書くセグメント木から ACL へ移行する

Table of Contents

はじめに

木マスター養成講座で普通のセグメント木ををソラでかけるようになったあとに、「遅延評価セグメント木をソラで書きたいあなたに」 (以降ソラ版と呼ぶ)を読んで遅延評価セグメント木をソラでかけるようになった。

とはいえ、ソラで書くの面倒なので ACL に移行するために「AtCoder LibraryのLazy Segtreeの使い方」を読んでみたものの、lazy 配列の使い方がソラ版と違っていて対応関係がよくわからなかったのでまとめる。

主な違いまとめ

ACL 版を使うときの S, op, e, F, mapping, composition はそれぞれ何を表しているのか

例

AtCoder Library Practice Contest, K - Range Affine Range Sum を見ながら何を書けばいいのかを確認していく

全体像は以下のよう

struct S {
    mint val;  // data node の値
    int size;  // data node が被覆している範囲
};

// 2つの子の data node から値を更新する
S op(S a, S b) {
    return {a.val + b.val, a.size + b.size};
}

// data node の単位元
S e() {
    return {0, 1};
}

// 区間適用する操作を表すデータ
// 末端の値に対して x <- b * x + c というデータの更新をするので b, c が必要
struct F {
    mint b, c;
};

// x <- b * x + c という操作の本体
S mapping(F f, S x) {
    return {
        f.b * x.val + f.c * x.size,
        x.size,
    };
}

F composition(F f, F g) {
    return {
        f.b * g.b,
        f.b * g.c + f.c,
    };
}

F id() {
    return {1, 0};
}

void solve() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, q;
    cin >> n >> q;

    lazy_segtree<S, op, e, F, mapping, composition, id> seg(n);

    rep(i, n) {
        ll a;
        cin >> a;
        seg.set(i, {a, 1});
    }

    rep(i, q) {
        int t;
        cin >> t;
        if (t == 0) {
            int l, r;
            ll b, c;
            cin >> l >> r >> b >> c;
            seg.apply(l, r, {b, c});
        } else {
            int l, r;
            cin >> l >> r;
            cout << seg.prod(l, r).val.val() << endl;
        }
    }
}

詳細

struct S {
    mint val;  // data node の値
    int size;  // data node が被覆している範囲
};

S は data node の型で。値と被覆範囲を持っている。ソラ版では $l$, $r$ の値を使えたので lazy の内容を data node に適用する際に範囲情報を含めて値を更新できたが、ACL 版では mapping に範囲情報が渡されないのでどこか別の場所に範囲を持っておく必要がある。 F のほうに範囲をもたせるのは多分無理なので S にもたせる

S op(S a, S b) {
    return {a.val + b.val, a.size + b.size};
}

2つの子の data node から値を更新する。 前半の a.val + b.val は子の値を合算している。普通のセグメント木でもやっている操作。 後半の a.size + b.size は子のサイズを合算して更新対象の data node が被覆している範囲を計算している。

S e() {
    return {0, 1};
}

op(e(), x) = op(x, e()) = x が成り立つように e を定義する

struct F {
    mint b, c;
};

区間適用する操作を表すデータ。 末端の値に対して $a_i \leftarrow b \times a_i + c$ というデータの更新をするので $b$, $c$ が必要

S mapping(F f, S x) {
    return {
        f.b * x.val + f.c * x.size,
        x.size,
    };
}

x はセグメント木の data node のいずれかの値。 $x = a_i+\cdots + a_{l+\text{(node size)}}$ とすると、$i \in [l, l+\text{(node size)})$ に対して $a_i \leftarrow b \times a_i + c$ を適用したあとの値は

\begin{align} &\sum_{i=l}^{l+\text{(node size)}-1} b \times a_i + c \nonumber \\ &= b \left( \sum_{i=l}^{l+\text{(node size)}-1} a_i \right) + c \times \text{(node size)} \nonumber \\ &= b \times x + c \times (\text{node size}) \nonumber \end{align}

となる。

F composition(F f, F g) {
    return {
        f.b * g.b,
        f.b * g.c + f.c,
    };
}

lazy node の内容を更に合算する操作

$g(x) = b_g x + c_g$, $f(x) = b_f x + c_f$ とすると、

\begin{align} f(g(x)) &= b_f (b_g x + c_g) + c_f \nonumber \\ &= b_f b_g x + b_f c_g + c_f \nonumber \\ f(g(x)) &= f \circ g (x) = b_{fg} x + c_{fg} \nonumber \end{align}

where $b_{fg} = b_f b_g$, $c_{fg} = b_f c_g + c_f$

となる。

F id() {
    return {1, 0};
}

mapping(id(), x) = x が成り立つように id を定義する。

$b = 1, c = 0$ のとき $b \times a_i + c = 1 \times a_i + 0 = a_i$ なので恒等演算になる。