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

遅延セグ木(遅いけど万能)(Codon推奨)

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

class LazySegTree:
    def __init__(self, n):
        self.n = n
        self.size = 1
        while self.size < n:
            self.size <<= 1

        self.H = 0
        x = self.size
        while x > 0:
            self.H += 1
            x >>= 1

        INF = 10**30
        m = 2 * self.size

        self.minv = [INF] * m
        self.maxv = [-INF] * m
        self.sumv = [0] * m

        self.add = [0] * m
        self.setv = [0] * m
        self.has_set = [0] * m

        self.length = [0] * m
        for i in range(self.size, self.size + n):
            self.length[i] = 1
        for i in range(self.size - 1, 0, -1):
            self.length[i] = self.length[i << 1] + self.length[i << 1 | 1]

    def _apply_set(self, i, v):
        self.minv[i] = v
        self.maxv[i] = v
        self.sumv[i] = v * self.length[i]
        if i < self.size:
            self.setv[i] = v
            self.add[i] = 0
            self.has_set[i] = 1

    def _apply_add(self, i, v):
        self.minv[i] += v
        self.maxv[i] += v
        self.sumv[i] += v * self.length[i]
        if i < self.size:
            if self.has_set[i]:
                self.setv[i] += v
            else:
                self.add[i] += v

    def _push(self, i):
        for s in range(self.H, 0, -1):
            p = i >> s
            if self.has_set[p]:
                v = self.setv[p]
                lc = p << 1
                self._apply_set(lc, v)
                self._apply_set(lc | 1, v)
                self.has_set[p] = 0
            if self.add[p] != 0:
                v = self.add[p]
                lc = p << 1
                self._apply_add(lc, v)
                self._apply_add(lc | 1, v)
                self.add[p] = 0

    def _recalc(self, i):
        while i > 1:
            i >>= 1
            if self.has_set[i]:
                v = self.setv[i]
                self.minv[i] = v
                self.maxv[i] = v
                self.sumv[i] = v * self.length[i]
            else:
                l = i << 1
                self.minv[i] = min(self.minv[l], self.minv[l | 1]) + self.add[i]
                self.maxv[i] = max(self.maxv[l], self.maxv[l | 1]) + self.add[i]
                self.sumv[i] = self.sumv[l] + self.sumv[l | 1] + self.add[i] * self.length[i]

    # ---- API ----

    # [l, r], 1-indexed
    def range_add(self, l, r, v):
        l += self.size - 1
        r += self.size - 1
        self._push(l)
        self._push(r)
        l0 = l
        r0 = r
        while l <= r:
            if (l & 1) == 1:
                self._apply_add(l, v)
                l += 1
            if (r & 1) == 0:
                self._apply_add(r, v)
                r -= 1
            l >>= 1
            r >>= 1
        self._recalc(l0)
        self._recalc(r0)

    def range_set(self, l, r, v):
        l += self.size - 1
        r += self.size - 1
        self._push(l)
        self._push(r)
        l0 = l
        r0 = r
        while l <= r:
            if (l & 1) == 1:
                self._apply_set(l, v)
                l += 1
            if (r & 1) == 0:
                self._apply_set(r, v)
                r -= 1
            l >>= 1
            r >>= 1
        self._recalc(l0)
        self._recalc(r0)

    def range_min(self, l, r):
        l += self.size - 1
        r += self.size - 1
        self._push(l)
        self._push(r)
        res = 10**30
        while l <= r:
            if (l & 1) == 1:
                if self.minv[l] < res:
                    res = self.minv[l]
                l += 1
            if (r & 1) == 0:
                if self.minv[r] < res:
                    res = self.minv[r]
                r -= 1
            l >>= 1
            r >>= 1
        return res

    def range_max(self, l, r):
        l += self.size - 1
        r += self.size - 1
        self._push(l)
        self._push(r)
        res = -10**30
        while l <= r:
            if (l & 1) == 1:
                if self.maxv[l] > res:
                    res = self.maxv[l]
                l += 1
            if (r & 1) == 0:
                if self.maxv[r] > res:
                    res = self.maxv[r]
                r -= 1
            l >>= 1
            r >>= 1
        return res

    def range_sum(self, l, r):
        l += self.size - 1
        r += self.size - 1
        self._push(l)
        self._push(r)
        res = 0
        while l <= r:
            if (l & 1) == 1:
                res += self.sumv[l]
                l += 1
            if (r & 1) == 0:
                res += self.sumv[r]
                r -= 1
            l >>= 1
            r >>= 1
        return res

    def point_set(self, pos, v):
        i = pos + self.size - 1
        self._push(i)
        self.minv[i] = v
        self.maxv[i] = v
        self.sumv[i] = v
        self.add[i] = 0
        self.has_set[i] = 0
        self._recalc(i)

    def point_get(self, pos):
        i = pos + self.size - 1
        self._push(i)
        return self.sumv[i]