遅延評価セグメント木の練習！

問題へのリンク

問題概要

長さ $N$ の数列 $a_{0}, a_{1}, \dots, a_{N-1}$ が与えられる。この数列に対して、次の $Q$ 回のクエリに答えよ。

クエリタイプ 1 ( $l, r, b, c$ )：数列の区間 $\lbrack l, r)$ 内の各要素の値を $b$ 倍して $c$ を足せ
クエリタイプ 2 ( $l, r$ )：数列の区間 $\lbrack l, r)$ 内の要素の総和を 998244353 で割った余りを答えよ

制約

$1 \le N, Q \le 5 \times 10^{5}$

解法：遅延評価セグメント木

数列の区間に対して「更新」と「取得」の両方がある場合は遅延評価セグメント木が有効になる可能性を疑います。遅延評価セグメント木の原理については、次の記事に詳しく書いてあります。

betrue12.hateblo.jp

遅延評価セグメント木では、次の要素を定義します。

数列の区間についての欲しい値 (総和など) を表す構造体 (ここでは `Node` とします)

2 つの構造体に対する二項演算を表す関数 (関数 Node op(Node, Node))
その二項演算 op における単位元を返す関数 (関数 Node e())

数列の区間の各要素に対する特定の作用を表す構造体 (ここでは `Act` とします)

構造体 Node に対する作用を表す関数 (関数 Node mapping(Act, Node))
作用と作用の合成関数 (関数 Act composition(Act, Act))
作用の単位元を返す関数 (関数 id())

今回の問題に対する遅延評価セグメント木の設計

今回の問題に対して、これらの要素を定義していきましょう。なお、区間の総和を取得したい問題において、区間の長さを Node に持たせるのは、セグメント木典型ですね。

区間の総和を表す構造体 `Node`

数列の区間 $\lbrack l, r)$ の各要素の総和を表す値 val と、区間の長さを表す siz (= $r - l$ ) の組を持たせます。

和の二項演算

Node op(Node x, Node y) { return Node(x.val + y.val, x.siz + y.siz); }

2 つの区間についての和をとるので、val と siz について、ともに和をとります。

単位元

Node e() { return Node(0, 0); }

作用を表す構造体 `Act`

$b$ 倍して $c$ を足すことを表す 2 つの値 $b, c$ を持たせます。

作用関数

Node mapping(Act f, Node x) {
    return Node(f.b * x.val + f.c * x.siz, x.siz);
}

まず、区間の各要素を $b$ 倍すると、総和も $b$ 倍になります。さらに各要素に $c$ を足すと、区間の長さは siz であることから、総和は $c * siz$ だけ増加します。よって、作用関数は上のように書けます。

作用の合成関数

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

要素 $x$ に対して、f.b 倍して f.c を足したあと、g.b 倍して g.c を足すという操作が、全体で見て何倍して何を足すのかを求めます。

g.b((f.b)x + f.c) + g.c = (g.b f.b)x + g.b f.c + g.c

なので、作用の合成関数は上のように書けます。

遅延評価セグメント木が使えるか

最後に、この問題に対して、遅延評価セグメント木が使えるかどうかを確認しましょう。遅延評価セグメント木が使えるためには、「作用したものと作用したものの総和」と「総和をとったものに作用したもの」が一致する必要があります。

数学的に書けば、構造体 Node の要素 $x, y$ に二項演算を適用したものを $x + y$ と書くことにして、構造体 Node の要素 $x$ へ作用したものを $f(x)$ と書くことにしたとき、

$f(x) + f(y) = f(x + y)$

が成り立つことが必要です。要素 $x, y$ の表す区間の長さを $l_{x}, l_{y}$ とすると、

$f(x) + f(y) = (bx + c l_{x}) + (by + c l_{y}) = b(x + y) + c(l_{x} + l_{y})$
$f(x + y) = b(x + y) + c(l_{x} + l_{y})$

ですので、確かに成り立っています。

よって、この遅延評価セグメント木を用いて、この問題は $O(N + Q \log N)$ の計算量で解けます。

コード

ACL を用いた実装

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

/* セグメント木のための構造体と、二項演算関数 op と、単位元関数 e */
using mint = modint998244353;
struct Node {
    mint val;
    long long siz;
    Node(mint v = 0, long long s = 0) : val(v), siz(s) {}
};

// 二項演算
Node op(Node x, Node y) { return Node(x.val + y.val, x.siz + y.siz); }

// 単位元
Node e() { return Node(0, 0); }

/* 遅延評価のための構造体と、作用関数 mapping と、作用の合成関数 composition と、単位元関数 id() */
struct Act {
    mint b, c;
    Act(mint b = 0, mint c = 0) : b(b), c(c) {}
};

// 作用関数
Node mapping(Act f, Node x) {
    return Node(f.b * x.val + f.c * x.siz, x.siz);
}

// 作用の合成関数：g.b((f.b)x + f.c) + g.c = (g.b f.b)x + g.b f.c + g.c
Act composition(Act g, Act f) {
    return Act(g.b * f.b, g.b * f.c + g.c);
}

// 作用の単位元
Act id() { return Act(1, 0); }

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<Node> a(N);
    for (int i = 0; i < N; ++i) {
        int x;
        cin >> x;
        a[i] = Node(x, 1);
    }
    
    // 遅延評価セグメント木のセットアップ
    lazy_segtree<Node, op, e, Act, mapping, composition, id> seg(a);
    
    // クエリ処理
    while (Q--) {
        int t;
        cin >> t;
        if (t == 0) {
            int l, r, c, d;
            cin >> l >> r >> c >> d;
            seg.apply(l, r, Act(c, d));
        } else {
            int l, r;
            cin >> l >> r;
            Node res = seg.prod(l, r);
            cout << res.val.val() << endl;
        }
    }
}

自前ライブラリでも AC

自前の LazySegmentTree と modint でも通します。

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

// Lazy Segment Tree
template<class Monoid, class Action> struct LazySegmentTree {
    // various function types
    using FuncMonoid = function<Monoid(Monoid, Monoid)>;
    using FuncAction = function<Monoid(Action, Monoid)>;
    using FuncComposition = function<Action(Action, Action)>;

    // core member
    int N;
    FuncMonoid OP;
    FuncAction ACT;
    FuncComposition COMP;
    Monoid IDENTITY_MONOID;
    Action IDENTITY_ACTION;
    
    // inner data
    int log, offset;
    vector<Monoid> dat;
    vector<Action> lazy;
    
    // constructor
    LazySegmentTree() {}
    LazySegmentTree(int n, const FuncMonoid op, const FuncAction act, const FuncComposition comp,
                    const Monoid &identity_monoid, const Action &identity_action) {
        init(n, op, act, comp, identity_monoid, identity_action);
    }
    LazySegmentTree(const vector<Monoid> &v,
                    const FuncMonoid op, const FuncAction act, const FuncComposition comp,
                    const Monoid &identity_monoid, const Action &identity_action) {
        init(v, op, act, comp, identity_monoid, identity_action);
    }
    void init(int n, const FuncMonoid op, const FuncAction act, const FuncComposition comp,
              const Monoid &identity_monoid, const Action &identity_action) {
        N = n, OP = op, ACT = act, COMP = comp;
        IDENTITY_MONOID = identity_monoid, IDENTITY_ACTION = identity_action;
        log = 0, offset = 1;
        while (offset < N) ++log, offset <<= 1;
        dat.assign(offset * 2, IDENTITY_MONOID);
        lazy.assign(offset * 2, IDENTITY_ACTION);
    }
    void init(const vector<Monoid> &v,
              const FuncMonoid op, const FuncAction act, const FuncComposition comp,
              const Monoid &identity_monoid, const Action &identity_action) {
        init((int)v.size(), op, act, comp, identity_monoid, identity_action);
        build(v);
    }
    void build(const vector<Monoid> &v) {
        assert(N == (int)v.size());
        for (int i = 0; i < N; ++i) dat[i + offset] = v[i];
        for (int k = offset - 1; k > 0; --k) pull_dat(k);
    }
    int size() const {
        return N;
    }
    
    // basic functions for lazy segment tree
    void pull_dat(int k) {
        dat[k] = OP(dat[k * 2], dat[k * 2 + 1]);
    }
    void apply_lazy(int k, const Action &f) {
        dat[k] = ACT(f, dat[k]);
        if (k < offset) lazy[k] = COMP(f, lazy[k]);
    }
    void push_lazy(int k) {
        apply_lazy(k * 2, lazy[k]);
        apply_lazy(k * 2 + 1, lazy[k]);
        lazy[k] = IDENTITY_ACTION;
    }
    void pull_dat_deep(int k) {
        for (int h = 1; h <= log; ++h) pull_dat(k >> h);
    }
    void push_lazy_deep(int k) {
        for (int h = log; h >= 1; --h) push_lazy(k >> h);
    }
    
    // setter and getter, update A[i], i is 0-indexed, O(log N)
    void set(int i, const Monoid &v) {
        assert(0 <= i && i < N);
        int k = i + offset;
        push_lazy_deep(k);
        dat[k] = v;
        pull_dat_deep(k);
    }
    Monoid get(int i) {
        assert(0 <= i && i < N);
        int k = i + offset;
        push_lazy_deep(k);
        return dat[k];
    }
    Monoid operator [] (int i) {
        return get(i);
    }
    
    // apply f for index i
    void apply(int i, const Action &f) {
        assert(0 <= i && i < N);
        int k = i + offset;
        push_lazy_deep(k);
        dat[k] = ACT(f, dat[k]);
        pull_dat_deep(k);
    }
    // apply f for interval [l, r)
    void apply(int l, int r, const Action &f) {
        assert(0 <= l && l <= r && r <= N);
        if (l == r) return;
        l += offset, r += offset;
        for (int h = log; h >= 1; --h) {
            if (((l >> h) << h) != l) push_lazy(l >> h);
            if (((r >> h) << h) != r) push_lazy((r - 1) >> h);
        }
        int original_l = l, original_r = r;
        for (; l < r; l >>= 1, r >>= 1) {
            if (l & 1) apply_lazy(l++, f);
            if (r & 1) apply_lazy(--r, f);
        }
        l = original_l, r = original_r;
        for (int h = 1; h <= log; ++h) {
            if (((l >> h) << h) != l) pull_dat(l >> h);
            if (((r >> h) << h) != r) pull_dat((r - 1) >> h);
        }
    }
    
    // get prod of interval [l, r)
    Monoid prod(int l, int r) {
        assert(0 <= l && l <= r && r <= N);
        if (l == r) return IDENTITY_MONOID;
        l += offset, r += offset;
        for (int h = log; h >= 1; --h) {
            if (((l >> h) << h) != l) push_lazy(l >> h);
            if (((r >> h) << h) != r) push_lazy(r >> h);
        }
        Monoid val_left = IDENTITY_MONOID, val_right = IDENTITY_MONOID;
        for (; l < r; l >>= 1, r >>= 1) {
            if (l & 1) val_left = OP(val_left, dat[l++]);
            if (r & 1) val_right = OP(dat[--r], val_right);
        }
        return OP(val_left, val_right);
    }
    Monoid all_prod() {
        return dat[1];
    }
    
    // get max r that f(get(l, r)) = True (0-indexed), O(log N)
    // f(IDENTITY) need to be True
    int max_right(const function<bool(Monoid)> f, int l = 0) {
        if (l == N) return N;
        l += offset;
        push_lazy_deep(l);
        Monoid sum = IDENTITY_MONOID;
        do {
            while (l % 2 == 0) l >>= 1;
            if (!f(OP(sum, dat[l]))) {
                while (l < offset) {
                    push_lazy(l);
                    l = l * 2;
                    if (f(OP(sum, dat[l]))) {
                        sum = OP(sum, dat[l]);
                        ++l;
                    }
                }
                return l - offset;
            }
            sum = OP(sum, dat[l]);
            ++l;
        } while ((l & -l) != l);  // stop if l = 2^e
        return N;
    }

    // get min l that f(get(l, r)) = True (0-indexed), O(log N)
    // f(IDENTITY) need to be True
    int min_left(const function<bool(Monoid)> f, int r = -1) {
        if (r == 0) return 0;
        if (r == -1) r = N;
        r += offset;
        push_lazy_deep(r - 1);
        Monoid sum = IDENTITY_MONOID;
        do {
            --r;
            while (r > 1 && (r % 2)) r >>= 1;
            if (!f(OP(dat[r], sum))) {
                while (r < offset) {
                    push_lazy(r);
                    r = r * 2 + 1;
                    if (f(OP(dat[r], sum))) {
                        sum = OP(dat[r], sum);
                        --r;
                    }
                }
                return r + 1 - offset;
            }
            sum = OP(dat[r], sum);
        } while ((r & -r) != r);
        return 0;
    }
    
    // debug stream
    friend ostream& operator << (ostream &s, LazySegmentTree seg) {
        for (int i = 0; i < (int)seg.size(); ++i) {
            s << seg[i];
            if (i != (int)seg.size() - 1) s << " ";
        }
        return s;
    }
    
    // dump
    void dump() {
        for (int i = 0; i <= log; ++i) {
            for (int j = (1 << i); j < (1 << (i + 1)); ++j) {
                cout << "{" << dat[j] << "," << lazy[j] << "} ";
            }
            cout << endl;
        }
    }
};

// modint
template<int MOD> struct Fp {
    // inner value
    long long val;
    
    // constructor
    constexpr Fp() : val(0) { }
    constexpr Fp(long long v) : val(v % MOD) {
        if (val < 0) val += MOD;
    }
    
    // getter
    constexpr long long get() const {
        return val;
    }
    constexpr int get_mod() const {
        return MOD;
    }
    
    // comparison operators
    constexpr bool operator == (const Fp &r) const {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp &r) const {
        return this->val != r.val;
    }
    
    // arithmetic operators
    constexpr Fp& operator += (const Fp &r) {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp &r) {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp &r) {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp &r) {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp operator + () const { return Fp(*this); }
    constexpr Fp operator - () const { return Fp(0) - Fp(*this); }
    constexpr Fp operator + (const Fp &r) const { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp &r) const { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp &r) const { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp &r) const { return Fp(*this) /= r; }
    
    // other operators
    constexpr Fp& operator ++ () {
        ++val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -- () {
        if (val == 0) val += MOD;
        --val;
        return *this;
    }
    constexpr Fp operator ++ (int) {
        Fp res = *this;
        ++*this;
        return res;
    }
    constexpr Fp operator -- (int) {
        Fp res = *this;
        --*this;
        return res;
    }
    friend constexpr istream& operator >> (istream &is, Fp<MOD> &x) {
        is >> x.val;
        x.val %= MOD;
        if (x.val < 0) x.val += MOD;
        return is;
    }
    friend constexpr ostream& operator << (ostream &os, const Fp<MOD> &x) {
        return os << x.val;
    }
    
    // other functions
    constexpr Fp pow(long long n) const {
        Fp res(1), mul(*this);
        while (n > 0) {
            if (n & 1) res *= mul;
            mul *= mul;
            n >>= 1;
        }
        return res;
    }
    constexpr Fp inv() const {
        Fp res(1), div(*this);
        return res / div;
    }
    friend constexpr Fp<MOD> pow(const Fp<MOD> &r, long long n) {
        return r.pow(n);
    }
    friend constexpr Fp<MOD> inv(const Fp<MOD> &r) {
        return r.inv();
    }
};

/* セグメント木のための構造体と、二項演算関数 op と、単位元関数 e */
using mint = Fp<998244353>;
struct Node {
    mint val;
    long long siz;
    Node(mint v = 0, long long s = 0) : val(v), siz(s) {}
};

// 二項演算
Node op(Node x, Node y) { return Node(x.val + y.val, x.siz + y.siz); }

// 単位元
const Node identity_monoid = Node(0, 0);

/* 遅延評価のための構造体と、作用関数 mapping と、作用の合成関数 composition と、単位元関数 id() */
struct Act {
    mint b, c;
    Act(mint b = 0, mint c = 0) : b(b), c(c) {}
};

// 作用関数
Node mapping(Act f, Node x) {
    return Node(f.b * x.val + f.c * x.siz, x.siz);
}

// 作用の合成関数：g.b((f.b)x + f.c) + g.c = (g.b f.b)x + g.b f.c + g.c
Act composition(Act g, Act f) {
    return Act(g.b * f.b, g.b * f.c + g.c);
}

// 作用の単位元
const Act identity_action = Act(1, 0);

int main() {
    // 入力
    int N, Q;
    cin >> N >> Q;
    vector<Node> a(N);
    for (int i = 0; i < N; ++i) {
        int x;
        cin >> x;
        a[i] = Node(x, 1);
    }
    
    // 遅延評価セグメント木のセットアップ
    LazySegmentTree<Node, Act> seg(a, op, mapping, composition,
                                   identity_monoid, identity_action);
    
    // クエリ処理
    while (Q--) {
        int t;
        cin >> t;
        if (t == 0) {
            int l, r, c, d;
            cin >> l >> r >> c >> d;
            seg.apply(l, r, Act(c, d));
        } else {
            int l, r;
            cin >> l >> r;
            Node res = seg.prod(l, r);
            cout << res.val << endl;
        }
    }
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder Library Practice Contest K - Range Affine Range Sum

問題概要

制約

解法：遅延評価セグメント木

数列の区間についての欲しい値 (総和など) を表す構造体 (ここでは `Node` とします)

数列の区間の各要素に対する特定の作用を表す構造体 (ここでは `Act` とします)

今回の問題に対する遅延評価セグメント木の設計

区間の総和を表す構造体 `Node`

和の二項演算

単位元

作用を表す構造体 `Act`

作用関数

作用の合成関数

遅延評価セグメント木が使えるか

コード

ACL を用いた実装

自前ライブラリでも AC

問題概要

制約

解法：遅延評価セグメント木

数列の区間についての欲しい値 (総和など) を表す構造体 (ここでは Node とします)

数列の区間の各要素に対する特定の作用を表す構造体 (ここでは Act とします)

今回の問題に対する遅延評価セグメント木の設計

区間の総和を表す構造体 Node

和の二項演算

単位元

作用を表す構造体 Act

作用関数

作用の合成関数

遅延評価セグメント木が使えるか

コード

ACL を用いた実装

自前ライブラリでも AC

数列の区間についての欲しい値 (総和など) を表す構造体 (ここでは `Node` とします)

数列の区間の各要素に対する特定の作用を表す構造体 (ここでは `Act` とします)

区間の総和を表す構造体 `Node`

作用を表す構造体 `Act`