抽象化遅延セグメント木の使い方 - プログラミング雑記帳

自分用に、抽象化した遅延セグメント木の使い方をとりあえず書いておきます。ここでは実際に幾つかの問題を解くことで使い方を示します。

抽象化遅延セグメント木の実装
Range Minimum Query and Range Update Query(RMQ and RUQ)
- コード
Range Affine Range Sum
- コード

抽象化遅延セグメント木の実装

まずはじめに、抽象化遅延セグメント木のコードを置いておきます。

template <typename T, typename E>
struct LazySegmentTree { // 0-indexed
    using F = function<T(T, T)>;
    using G = function<T(T, E)>;
    using H = function<E(E, E)>;
    
    int n, height;
    vector<T> tree;
    vector<E> lazy;
    const F f;
    const G g;
    const H h;
    const T ti;
    const E ei;

    LazySegmentTree() {}
    LazySegmentTree(F f, G g, H h, T ti, E ei) :f(f), g(g), h(h), ti(ti), ei(ei) {}

    void init(int n_) {
        n = 1, height = 0;
        while (n < n_) n *= 2, ++height;
        tree.assign(2 * n, ti);
        lazy.assign(2 * n, ei);
    }

    void build(const vector<T>& v) {
        int n_ = v.size();
        init(n_);
        for (int i = 0; i < n_; ++i) tree[n + i] = v[i];
        for (int i = n - 1; i > 0; --i) tree[i] = f(tree[2 * i], tree[2 * i + 1]);
    }

    inline T reflect(int k) { return (lazy[k] == ei?tree[k]:g(tree[k], lazy[k])); }

    inline void eval(int k) {
        if (lazy[k] == ei) return;
        lazy[2 * k] = h(lazy[2 * k], lazy[k]);
        lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
        tree[k] = reflect(k);
        lazy[k] = ei;
    }

    inline void thrust(int k) { for (int i = height; i > 0; --i) eval(k >> i); }

    inline void recalc(int k) {
        while (k >>= 1)
            tree[k] = f(reflect(2 * k), reflect(2 * k + 1));
    }
    
    void update(int s, int t, const E& x) { // [l, r)
        s += n, t += n;
        thrust(s), thrust(t - 1);
        int l = s, r = t;
        while (l < r) {
            if (l & 1) lazy[l] = h(lazy[l], x), ++l;
            if (r & 1) --r, lazy[r] = h(lazy[r], x);
            l >>= 1, r >>= 1;
        }
        recalc(s), recalc(t - 1);
    }

    T find(int s, int t) { // [l, r)
        s += n, t += n;
        thrust(s), thrust(t - 1);
        int l = s, r = t;
        T ll = ti, rr = ti;
        while (l < r) {
            if (l & 1) ll = f(ll, reflect(l++));
            if (r & 1) rr = f(rr, reflect(--r));
            l >>= 1, r >>= 1;
        }
        return f(ll, rr);
    }

    T at(int i) { return find(i, i + 1); }
};

詳しくは説明しませんが、 $T$ は取得クエリに必要な情報を表す型、 $E$ は更新クエリに必要な情報を表す型、 $f$ は $T$ と $T$ をマージする関数、 $g$ は $T$ に $E$ を作用させる関数、 $h$ は $E$ と $E$ をマージする関数、 $ti$ は $T$ の $f$ に関する単位元、 $ei$ は $E$ の $h$ に関する単位元です。

Range Minimum Query and Range Update Query(RMQ and RUQ)

これはAOJのこの問題です。一応、概要を書いておくと、

数列 $a_0, a_1, a_2, \cdots, a_{n-1}$ に対し、

$a_s, a_{s+1}, \cdots, a_t$ を $x$ に変更する。
$a_s, a_{s+1}, \cdots, a_t$ の最小値を出力する。

この問題において取得クエリは $find(s, t)$ なので、取得クエリに必要な情報は最小値のみです。したがって $T$ はintとなります。また、更新クエリは $update(s, t, x)$ なので、更新クエリに必要な情報は $x$ のみです。したがって $E$ もintとなります。

次に、関数 $f$ を考えます。 $A_{i,j} = min(a_i, a_{i+1}, \cdots, a_j)$ とすると、 $$A_{s,t} = min(A_{s,m}, A_{m+1,t}) \ (s \leq m \leq t)$$ と表せるので、 $$f(t1:T, t2:T) = min(t1, t2)$$ で良いことが分かります。

次に、関数 $g$ を考えます。 $update(s, t, x)$ のクエリのあとに、 $A_{s,t}$ がどのようになるかを考えると、明らかに $A_{s,t} = x$ なので、 $$g(t1:T, e1:E) = e1$$ で良いことが分かります。 \begin{eqnarray} g(t1:T, e1:E) = \begin{cases} t1 \ (e1 = ei) \\ e1 \ (otherwise) \end{cases} \end{eqnarray} とする方が良いかもしれません。

次に、関数 $h$ を考えます。 $update(s, t, x_1)$ のクエリのあとに、 $update(s, t, x_2)$ のクエリが来たとき、持つ必要がある情報は明らかに $x_2$ だけで十分なので、 $$h(e1:E, e2:E) = e2$$ で良いことが分かります(これは実装において、第二引数をあとに適用するようにしているためです)。 \begin{eqnarray} h(e1:E, e2:E) = \begin{cases} e1 \ (e2 = ei) \\ e2 \ (otherwise) \end{cases} \end{eqnarray} とする方が良いかもしれません。

最後に単位元 $ti$ と $ei$ を考えます。 intの $f$ に関する単位元は $INF$ としておくと都合が良いので $ti = INF$ などとします。しかし、今回の場合は $2^{31}-1$ で初期化されているので $$ti = 2^{31}-1$$ とします。 $h$ は代入演算とみなすことができ、更新クエリで与えられないような値に設定しておくとよいので $$ei = -1$$ などとします。

コード

int n, q; cin >> n >> q;
auto f = [](int a, int b) { return min(a, b); };
auto g = [](int a, int b) { return b; };
LazySegmentTree<int, int> seg(f, g, g, INT_MAX, -1);
seg.init(n);
while (q--) {
    int com, s, t, x; cin >> com >> s >> t;
    if (com) cout << seg.find(s, t + 1) << endl;
    else cin >> x, seg.update(s, t + 1, x);
}

https://onlinejudge.u-aizu.ac.jp/solutions/problem/DSL_2_F/review/4027109/S_ZK/C++14

Range Affine Range Sum

これはAtCoderのこの問題です。一応、概要を書いておくと、

数列 $a_0, a_1, a_2, \cdots, a_{n-1}$ に対し、

$a_l, a_{l+1}, \cdots, a_{r-1}$ のそれぞれを $a_i = ba_i+c$ にする。
$\displaystyle \left(\sum_{i=l}^{r-1} a_i\right) mod \ 998244353$ を出力する。

この問題では $T$ や $E$ に区間の長さを持たせることで遅延セグメント木を使って解くことができます。また、クエリの区間が半開区間なのに注意します。

以下では $mod \ 998244353$ と明記しませんが、実装時には適切に書く必要があります。

$T$ を区間の合計値と区間の長さを保持する型とします。 $E$ を $a_i = ba_i+c$ の $b$ と $c$ を保持する型とします。

ここで、関数 $f$ を考えます。 $A_{[i,j)}^{sum}$ を $a_i, a_{i+1}, \cdots, a_{j-1}$ の合計、 $A_{[i,j)}^{len}$ を区間の長さとし、 $A_{i,j} = (A_{[i,j)}^{sum}, \ A_{[i,j)}^{len})$ とすると、 $$f(A_{l,m}:T, A_{m,r}:T) = (A_{[l,m)}^{sum}+A_{[m,r)}^{sum}, \ A_{[l,m)}^{len}+A_{[m,r)}^{len})$$ となります。

次に、関数 $g$ を考えます。 $a_l, a_{l+1}, \cdots, a_{r-1}$ は更新クエリのあと、 $$ba_l+c, \ ba_{l+1}+c, \cdots, \ ba_{r-1}+c$$ となるため、 $e_1 = (e_1^b, e_1^c) = (b, c)$ とすると、 $$g(A_{l,r}:T, e_1:E) = (e_1^bA_{[l,r)}^{sum}+e_1^cA_{[l,r)}^{len}, \ A_{[l,r)}^{len})$$ となります。

次に、関数 $h$ を考えます。 $e_1 = (b_1, c_1), e_2 = (b_2, c_2)$ とすると、 $a_l, a_{l+1}, \cdots, a_{r-1}$ は、更新クエリ $e_1$ のあとに $$b_1a_l+c_1, \ b_1a_{l+1}+c_1, \cdots, \ b_1a_{r-1}+c_1$$ となり、この次に更新クエリ $e_2$ が来ると、 $$b_2b_1a_l+b_2c_1+c_2, \ b_2b_1a_{l+1}+b_2c_1+c_2, \cdots, \ b_2b_1a_{r-1}+b_2c_1+c_2$$ となります。したがって、 $$h(e_1:E, e_2:E) = (e_2^be_1^b, \ e_2^be_1^c+e_2^c)$$ となります。

最後に単位元 $ti$ と $ei$ を考えます。 $T$ の $f$ に関する単位元は、 $T$ が合計値と長さを持つため、 $ti = (0, 0)$ となります。 $E$ の $h$ に関する単位元は、 $h(e_1:E, e_2:E) = e_1$ を満たすような $e_2$ となるため、 $ei = (1, 0)$ となります。

コード

using mint = ModInt<int, 998244353>;
int N, Q; cin >> N >> Q;
vector<pair<mint, mint> > A(N);
for (auto& [sum, len] : A) {
    cin >> sum;
    len = 1;
}
auto f = [](pair<mint, mint> a, pair<mint, mint> b) {
    return pair(a.first + b.first, a.second + b.second);
};
auto g = [](pair<mint, mint> a, pair<mint, mint> b) {
    return pair(b.first * a.first + b.second * a.second, a.second);
};
auto h = [](pair<mint, mint> a, pair<mint, mint> b) {
    return pair(b.first * a.first, b.first * a.second + b.second);
};
LazySegmentTree<pair<mint, mint>, pair<mint, mint> > seg(f, g, h, pair(0, 0), pair(1, 0));
seg.build(A);
while (Q--) {
    int f, l, r; cin >> f >> l >> r;
    if (f) cout << seg.find(l, r).first << '\n';
    else {
        mint b, c; cin >> b >> c;
        seg.update(l, r, pair(b, c));
    }
}

https://atcoder.jp/contests/practice2/submissions/18946367

ModInt は自分で用意してください...