< uruneaさんのAI生成コード置き場

遅延セグ木(モノイド設計)

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

class LazySegTree:
    def __init__(self, op, e, mapping, composition, id_v, v):
        n = len(v)
        log = 0
        while (1 << log) < n:
            log += 1
        size = 1 << log
        
        d = [e] * (2 * size)
        lz = [id_v] * size
        
        for i in range(n):
            d[size + i] = v[i]

        self._op = op
        self._e = e
        self._mapping = mapping
        self._composition = composition
        self._id = id_v
        
        self._n = n
        self._log = log
        self._size = size
        self._d = d
        self._lz = lz
        
        for i in range(self._size - 1, 0, -1):
            self._update(i)

    def _update(self, k):
        self._d[k] = self._op(self._d[2 * k], self._d[2 * k + 1])

    def _all_apply(self, k, f):
        self._d[k] = self._mapping(f, self._d[k])
        if k < self._size:
            self._lz[k] = self._composition(f, self._lz[k])

    def _push(self, k):
        if self._lz[k] != self._id:
            self._all_apply(2 * k, self._lz[k])
            self._all_apply(2 * k + 1, self._lz[k])
            self._lz[k] = self._id
    
    def set(self, p, x):
        p += self._size
        for i in range(self._log, 0, -1):
            self._push(p >> i)
        self._d[p] = x
        for i in range(1, self._log + 1):
            self._update(p >> i)

    def get(self, p):
        p += self._size
        for i in range(self._log, 0, -1):
            self._push(p >> i)
        return self._d[p]

    def prod(self, l, r):
        if l == r:
            return self._e
        l += self._size
        r += self._size
        for i in range(self._log, 0, -1):
            if ((l >> i) << i) != l: self._push(l >> i)
            if ((r >> i) << i) != r: self._push((r - 1) >> i)
        
        sml, smr = self._e, self._e
        while l < r:
            if l & 1:
                sml = self._op(sml, self._d[l])
                l += 1
            if r & 1:
                r -= 1
                smr = self._op(self._d[r], smr)
            l >>= 1
            r >>= 1
        return self._op(sml, smr)

    def all_prod(self):
        return self._d[1]

    def apply(self, l, r, f):
        if l == r:
            return
        l += self._size
        r += self._size
        for i in range(self._log, 0, -1):
            if ((l >> i) << i) != l: self._push(l >> i)
            if ((r >> i) << i) != r: self._push((r - 1) >> i)
        
        l2, r2 = l, r
        while l < r:
            if l & 1:
                self._all_apply(l, f)
                l += 1
            if r & 1:
                r -= 1
                self._all_apply(r, f)
            l >>= 1
            r >>= 1
        
        l, r = l2, r2
        for i in range(1, self._log + 1):
            if ((l >> i) << i) != l: self._update(l >> i)
            if ((r >> i) << i) != r: self._update((r - 1) >> i)

    def apply_point(self, p, f):
        p += self._size
        for i in range(self._log, 0, -1):
            self._push(p >> i)
        self._d[p] = self._mapping(f, self._d[p])
        for i in range(1, self._log + 1):
            self._update(p >> i)

mod = 998244353
nu = input().split()
N, Q = int(nu[0]), int(nu[1])

#右区間と左区間をマージ
def op(x, y):
    x_sum, x_sum2, x_leng = x
    y_sum, y_sum2, y_leng = y
    return ((x_sum+y_sum)%mod, (x_sum2+y_sum2)%mod, x_leng+y_leng)

#値の単位元
#op(x, e) = op(e, x) = x を満たす必要がある。
def e():
    return (0, 0, 0)

#区間に区間変更値fを付与
def mapping(f, S):
    s_sum, s_sum2, s_leng = S
    f %= mod
    new_sum = (s_sum + f * s_leng) % mod
    f2_leng = (f * f % mod) * s_leng % mod
    term2 = (2 * f % mod) * s_sum % mod
    new_sum2 = (s_sum2 + term2 + f2_leng) % mod
    return (new_sum, new_sum2, s_leng)

#元々溜まっていたgに新たに降ってきたfを合わせる
def composition(f, g):
    return (f + g)%mod

#何も操作しないときの状態(操作の単位元)
#mapping(id, x) = x と composition(f, id) = f を満たす必要がある。
#例えば区間乗算なら、1を掛けても何も変わらないのでid=1
#区間代入なら、入力で与えられない極端な数字6*10**18など
def id():
    return 0

aa = [(0, 0, 1) for i in range(N)]
seg = LazySegTree(op, (0, 0, 0), mapping, composition, 0, aa)