Day 041-Q3 — 遅延SegTree Affine(ax+b変換・区間和クエリ)

2026-05-24 赤色 Master / Phase 8+ ★★★★★★★★★ Lazy Segment Tree + Affine Transform

問題

長さ $N$ の数列 $A$ に対して以下のクエリを処理せよ:

  • 1 l r a b : $l \le i \le r$ の全要素に $A_i \leftarrow a \cdot A_i + b$ を適用。
  • 2 l r : $\sum_{i=l}^{r} A_i \pmod{998244353}$ を出力。

制約

$1 \le N, Q \le 2 \times 10^5$
$0 \le A_i, a, b < 998244353$
時間制限: 2sec / メモリ: 256MB

入出力例

入力例 1

5 3
1 2 3 4 5
1 2 4 2 1
2 1 5
1 1 5 0 7
2 1 5

出力例 1

25
35

概念図: Affine 変換の合成

f₁(x) = a₁x + b₁ 先に適用 f₂(x) = a₂x + b₂ 後に適用 合成: f₂∘f₁(x) = (a₂a₁)x + (a₂b₁+b₂) push_down 時の更新: ノードが lazy (a, b) を持つとき、子ノードに伝播する sum_new = a × sum_old + b × size ← サイズ情報が必要なので sz を各ノードに保持する

ヒント(段階的開示)

ヒント1: 方向性
区間 affine 変換 $f(x) = ax + b$ は合成可能な半群:$g \circ f(x) = (g_a \cdot a)x + (g_a \cdot b + g_b)$。遅延セグメント木の lazy 合成に使える。
ヒント2: アプローチ
  • ノードに「区間和」と「区間サイズ」を保持。
  • lazy に affine 変換 $(a, b)$ を保持(初期値 $(1, 0)$)。
  • push_down: $\text{sum}' = a \cdot \text{sum} + b \cdot \text{size}$
  • 合成: $(a_2, b_2) \circ (a_1, b_1) = (a_2 a_1, a_2 b_1 + b_2)$
ヒント3: 実装骨格
class LazySegTree:
    def _apply(self, i, a, b):
        self.sm[i] = (a * self.sm[i] + b * self.sz[i]) % MOD
        self.la[i] = a * self.la[i] % MOD
        self.lb[i] = (a * self.lb[i] + b) % MOD

    def _push(self, i):
        if self.la[i] != 1 or self.lb[i] != 0:
            self._apply(2*i, self.la[i], self.lb[i])
            self._apply(2*i+1, self.la[i], self.lb[i])
            self.la[i] = 1
            self.lb[i] = 0

模範解答 (Python)

import sys
input = sys.stdin.readline
MOD = 998244353

class LazySegTree:
    def __init__(self, A):
        n = len(A)
        self.n = 1
        while self.n < n: self.n <<= 1
        self.sm = [0] * (2 * self.n)
        self.sz = [0] * (2 * self.n)
        self.la = [1] * (2 * self.n)
        self.lb = [0] * (2 * self.n)
        for i, v in enumerate(A):
            self.sm[self.n + i] = v % MOD
            self.sz[self.n + i] = 1
        for i in range(self.n - 1, 0, -1):
            self.sm[i] = (self.sm[2*i] + self.sm[2*i+1]) % MOD
            self.sz[i] = self.sz[2*i] + self.sz[2*i+1]

    def _apply(self, i, a, b):
        self.sm[i] = (a * self.sm[i] + b * self.sz[i]) % MOD
        self.la[i] = a * self.la[i] % MOD
        self.lb[i] = (a * self.lb[i] + b) % MOD

    def _push(self, i):
        if self.la[i] != 1 or self.lb[i] != 0:
            self._apply(2*i, self.la[i], self.lb[i])
            self._apply(2*i+1, self.la[i], self.lb[i])
            self.la[i] = 1
            self.lb[i] = 0

    def update(self, l, r, a, b):  # [l, r] 1-indexed
        def rec(i, lo, hi):
            if r < lo or hi < l: return
            if l <= lo and hi <= r:
                self._apply(i, a, b); return
            self._push(i)
            mid = (lo + hi) >> 1
            rec(2*i, lo, mid)
            rec(2*i+1, mid+1, hi)
            self.sm[i] = (self.sm[2*i] + self.sm[2*i+1]) % MOD
        rec(1, 1, self.n)

    def query(self, l, r):  # [l, r] 1-indexed
        def rec(i, lo, hi):
            if r < lo or hi < l: return 0
            if l <= lo and hi <= r: return self.sm[i]
            self._push(i)
            mid = (lo + hi) >> 1
            return (rec(2*i, lo, mid) + rec(2*i+1, mid+1, hi)) % MOD
        return rec(1, 1, self.n)

def solve():
    N, Q = map(int, input().split())
    A = list(map(int, input().split()))
    seg = LazySegTree(A)
    out = []
    for _ in range(Q):
        q = list(map(int, input().split()))
        if q[0] == 1:
            _, l, r, a, b = q
            seg.update(l, r, a, b)
        else:
            _, l, r = q
            out.append(seg.query(l, r))
    print('\n'.join(map(str, out)))

solve()

Step-by-Step 解説

1Affine 変換の合成
$f_1(x) = a_1 x + b_1$, $f_2(x) = a_2 x + b_2$ の合成は $f_2(f_1(x)) = (a_2 a_1)x + (a_2 b_1 + b_2)$。これが遅延セグメント木の lazy 合成に対応する。
2push_down の正しさ
区間サイズが $s$ のとき、$a \cdot \text{sum} + b \cdot s$ が新しい sum。サイズ情報 sz を各ノードに保持しておく必要がある。
3update の流れ
根から下りて対象区間に _apply し、上りながら sum を再計算(sm[i] = sm[2i] + sm[2i+1])。
4query の流れ
根から下りる途中で _push を呼び、lazy を子に伝播してから正確な sum を集計する。
5単位元の設定
lazy の初期値は $(a=1, b=0)$(恒等写像)。_push で「変化なし」の判定に使う。

計算量

初期化: $O(N)$
区間 affine 更新: $O(\log N)$
区間和クエリ: $O(\log N)$
全体: $O((N + Q) \log N)$

よくあるミス

ミス原因正しい書き方
lazy の合成順序を逆に$f_2 \circ f_1$ の順新しい変換が「外側」
sz を管理しないpush_down で b*sz が必要葉以外も sz を正しく計算
push_down を葉でも呼ぶ葉は子がない葉のpushは不要(再帰の基底で自動回避)
MOD の取り忘れ大きな a, b の積各乗算後に % MOD

次のステップ

  • 発展問題: 区間 affine + 区間最大値(モノイドを拡張)
  • 応用: RGB 分解・行列型遅延伝播(2×2 行列の合成)
  • 類題: AtCoder Library の lazysegtree でアクション型を変えるパターン

自己評価