遅延伝播セグメント木の使い方，ACLPC: K – Range Affine Range Sum の解説

競プロの C++ ライブラリとして，AtCoder Library が提供されました．この中には遅延伝播セグメント木（lazysegtree）も入っています．
→ 公式告知，github

この記事は，「遅延じゃない普通のセグメント木もよく知らない…」「公式ドキュメントは数学用語が多くて読めない…」という方が，lazysegtree をブラックボックスとして使えるようになることを目標にしています．

例題として，AtCoder Library Practice Contest: K - Range Affine Range Sum を解説します．

「モノイド」「写像」等の用語は一切使っていない代わりに，厳密さを少し犠牲にしています．

2020/09/28 追記: 練習問題記事も書きました！ → 遅延セグ木に慣れる　AtCoder の練習問題 3 問

遅延伝播セグメント木とは

長さ $N$ の列 $(a_0, a_1,\dots, a_{N-1})$ が与えられます．以下では $0 \leq l \leq r \leq N$ とし，区間を全て半開区間 $[l, r) = \{l, l+1, \dots, r-1\}$ で書きます．また，連続部分列 $(a_l, a_{l+1}, \dots, a_{r-1})$ のことを $a_{[l, r)}$ と書きます．

遅延伝播セグメント木（遅延セグ木）は，任意の区間 $[l, r)$ に対し以下の $2$ つのクエリが処理できるデータ構造です:

遅延セグ木を使うためには，各区間 $[l, r)$ に対し，$a_{[l, r)}$ から必要な情報だけを上手く抽出する必要があります．この抽出した情報を，以下では $s_{[l, r)}$ と書きます．

$s_{[l, r)}$ はなんでもいいわけではなく，以下の条件を満たさないといけません:

「高速に」というのは，$O(1)$ 時間くらいを想定しています．これらの条件を満たすように上手く $s_{[l, r)}$ を構成するのが腕の見せどころというわけです．

更新クエリもなんでもいいわけではありませんが，少なくとも以下の条件を満たすなら大丈夫です¹:

処理したい更新クエリが条件 5 を満たしていない，つまり，$h$ が「同種の」クエリではなくなってしまう場合もあります．その場合，更新クエリの意味をより広く解釈すると上手くいくこともあるのですが，これは少し高度な話なのでこの記事では扱いません．

例題: ACLPC: K - Range Affine Range Sum

実際にどのような思考過程で遅延セグ木の問題を解くのか，例題で解説します．AtCoder Library の lazysegtree を使って実装します．

問題ページ: AtCoder Library Practice Contest: K - Range Affine Range Sum

$s_{[l, r)}$ の素朴な構成

まず，$a_{[l, r)}$ から抽出する情報 $s_{[l, r)}$ を構成し，上で説明した $3$ つの条件を満たすかどうかを確認します．

取得クエリは，$a_{[l, r)}$ の総和 $\sum_{i=l}^{r-1} a_i$ を答えよ，という形です．$s_{[l, r)}$ についての「条件 1」から，$s_{[l, r)} := \sum_{i=l}^{r-1} a_i$ とすれば良さそうに見えます．取得クエリが来たら，$s_{[l, r)}$ をそのまま答えればよいです．

こうすると条件 2 も満たされます．確かに $s_{[l, r)}$ は $s_{[l, r)} = s_{[l, m)} + s_{[m, r)}$ と計算できます．

条件 3 はどうでしょうか．更新クエリにより $a_{[l, r)}$ が $a'_{[l, r)} = (a'_l,\dots,a'_{r-1})$ に更新されたとしましょう．それと同時に，$s_{[l, r)} = \sum_{i=l}^{r-1} a_i$ が $s'_{[l, r)} = \sum_{i=l}^{r-1} a'_i$ に更新できないといけません．

この問題では $a'_i = b a_i + c$ となっています．実際に $s'_{[l, r)}$ を計算してみると，\begin{equation}\label{sum-update}s'_{[l, r)}= \sum_{i=l}^{r-1} (ba_i + c)= b s_{[l, r)} + c\under{(r-l)}\end{equation}です．

ここで，赤線の部分で困ってしまいます．$r-l$ は $a_{[l, r)}$ の長さを表しますが，この情報は $s_{[l, r)}$ を構成する時に捨ててしまったので，もう持っていません．つまり，$s_{[l, r)} := \sum_{i=l}^{r-1} a_i$ という定義では $s_{[l, r)}$ を更新できず，条件 3 を満たしません．

$s_{[l, r)}$ に区間長も持たせる

では，$s_{[l, r)}$ に $a_{[l, r)}$ の総和だけでなく，$a_{[l, r)}$ の長さ $r-l$ も持たせたらどうでしょうか？　つまり，\begin{equation}\label{s-def}s_{[l, r)}=(s_{[l, r)}^{\mathtt{sum}}, s_{[l, r)}^{\texttt{len}}):=\red{\bigg(\sum_{i=l}^{r-1} a_i,\ r-l\bigg)}\end{equation}と $2$ つの値を持つことにし，再度条件 1 〜 3 を確認しましょう．

情報が増えただけなので，もちろん条件 1 は満たしています．条件 2 も\begin{equation}\label{s-merge}s_{[l, r)}^{\texttt{sum}} = \red{s_{[l, m)}^{\texttt{sum}} + s_{[m, r)}^{\texttt{sum}}},\quad s_{[l, r)}^{\texttt{len}} = \red{s_{[l, m)}^{\texttt{len}} + s_{[m, r)}^{\texttt{len}}}\end{equation}とすればよいです．

問題は条件 3 でしたが，式 \eqref{sum-update} と「更新クエリで区間長が変更されない（当然！）」ことに注意すると，\begin{equation}\label{s-update}{s'}_{[l, r)}^{\texttt{sum}}= \red{b s_{[l, r)}^{\texttt{sum}} + c s_{[l, r)}^{\texttt{len}}},\quad{s'}_{[l, r)}^{\texttt{len}}= \red{s_{[l, r)}^{\texttt{len}}}\end{equation}と更新できます．よって条件 3 も満たされます！

更新クエリに関する条件

$s_{[l, r)}$ に関する条件の確認はできたので，更新クエリに関する条件 4 〜 5 を確認しましょう．

まず，$a_i$ の更新規則は $a_i \gets b a_i + c$ で，これは $a_j$ $(j \neq i)$ に依存していないので，条件 4 は満たしています．

次に条件 5 です．更新クエリ $g$ で $a'_i = b_g a_i + c_g$ と更新された後，更新クエリ $f$ で $a''_i = b_f a'_i + c_f$ と更新されるとします．代入して愚直に計算すると，$a''_i = (b_f b_g) a_i + (b_f c_g + c_f)$ となります．

つまり，更新クエリ $g,f$ を順に適用することは，\begin{equation}\label{composition}b := \red{b_f b_g},\qquad c := \red{b_f c_g + c_f}\end{equation}とした更新クエリを $1$ 回適用することと等価だと分かります．

これで 5 つ全ての条件の成立が確認できました！

AtCoder Library を使った実装

5 つの条件を確認する過程で出てきた式たちを使って実装します．

AtCoder Library を使って解く場合は，$2$ つの構造体 S, F と $5$ つの関数 op, e, mapping, composition, id を実装し，これらを「初期化 vector」とともに lazysegtree に渡します．以下，これについて説明します．

なお，この問題ではクエリの答えが $\mathrm{mod}\ \ 998244353$ で聞かれているので，最初に

using mint = modint998244353;

としておきます．

構造体の実装

構造体 S は，$a_{[l, r)}$ から抽出した情報 $s_{[l, r)}$ の型を表します．今回は，式 \eqref{s-def} にあるとおり $s_{[l, r)} = (s_{[l, r)}^{\texttt{sum}}, s_{[l, r)}^{\texttt{len}})$ としたので，以下のように書きます:

struct S {
  mint sum, len;
};

構造体 F は，$a_{[l, r)}$ に適用される更新クエリの情報を表す型です．今回は更新クエリとして（$l, r$ 以外に） $2$ つの整数 $b, c$ が与えられるので，以下のようにします:

struct F {
  mint b, c;
};

関数の実装

関数 op には，条件 2 を確認する際に得られた，$s_{[l, m)}, s_{[m, r)}$ から $s_{[l, r)}$ を計算する式を書きます．式 \eqref{s-merge} を実装すればよいです:

S op(S s_l, S s_r) {
  return {s_l.sum + s_r.sum, s_l.len + s_r.len};
}

関数 e には，長さ $0$ の区間に対応する情報 $s_{[i, i)}$ を書きます（これは $i$ に依存しません）．長さ $0$ の区間は総和も $0$ なので，$s_{[i, i)} = (s_{[i, i)}^{\texttt{sum}}, s_{[i, i)}^{\texttt{len}}) = (0, 0)$ です:

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

関数 mapping には，更新クエリによって $s_{[l, r)}$ がどのように変化するかを書きます．式 \eqref{s-update} を実装すればよいです:

S mapping(F f, S s) {
  return {f.b * s.sum + f.c * s.len, s.len};
}

関数 composition には，条件 5 を確認する際に得られた式を書きます．式 \eqref{composition} を実装すればよいです．クエリを適用する順番は $g,f$ ですが，引数は f, g の順で書くことに注意してください:

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

関数 id には「何もしない」更新クエリを書きます．今回は $(b, c) = (1, 0)$ が何もしないことに対応します:

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

初期化 vector の構成

この問題では整数列の初期値が与えられます．lazysegtree を作るときに，この情報も vector で渡してあげます．

ここで注意が必要なのは，初期値 $(a_0, a_1,\dots, a_{N-1})$ をそのまま渡すのではなく，長さ $1$ の各区間に対応する情報たち $(s_{[0, 1)}, s_{[1, 2)},\dots, s_{[N-1, N)})$ を渡すということです．今回は $s_{[i, i+1)} = (a_i, 1)$ です．

入力受け取りと lazysegtree の初期化も含めると，実装は以下のようになります:

int n, q; cin >> n >> q;
vector<S> v(n);
rep(i, n) {
  int a; cin >> a;
  v[i] = {a, 1};
}
lazy_segtree<S, op, e, F, mapping, composition, id> seg(v);

以上を踏まえて実装したコードがこちらです．→ AC コード (C++)

補足

• 上では（説明の都合もあり）例題を解く際に条件 1 から順に確認しましたが，実際には条件 4 〜 5 を確認してから条件 1 〜 3 を確認する方が，無駄がなくてよいと思います．
• 上の問題では「何もしない」更新クエリ $(b, c) = (1, 0)$ が来る可能性もあったのですが，問題によってはそのような更新クエリが決して来ないこともあります．実はそのような問題でも，（lazysegtree の性質上）何もしない更新クエリも来うるものと思って実装する必要があります．
• 何もしない更新クエリ id が $(b, c) = (1, 0)$ のようにきれいに書けない場合は，INF などの「ありえない値」が id に対応する，と自分で勝手に決めてしまい，関数 mapping, composition 内ではクエリを表す変数が「ありえない値」かどうかで場合分けをすれば上手くいきます．

練習問題

この記事ではカバーしきれなかった部分を補うために，練習問題解説記事を書きました．こちらにもぜひ挑戦してみてください．
→ 遅延セグ木に慣れる　AtCoder の練習問題 3 問

1. 2020/09/27 17:00 更新: 更新前は条件 4 が抜けていました．熨斗袋 (@noshi91) さんの指摘を受けて条件 4 を追加しました．ありがとうございます．