segment_lazy_quadratic_division

ソースコード

from titan_pylib.data_structures.segment_quadratic_division.segment_lazy_quadratic_division import SegmentLazyQuadraticDivision

view on github

展開済みコード

  1# from titan_pylib.data_structures.segment_quadratic_division.segment_lazy_quadratic_division import SegmentLazyQuadraticDivision
  2from typing import Union, Callable, TypeVar, Generic, Iterable
  3from functools import reduce
  4from itertools import chain
  5
  6T = TypeVar("T")
  7F = TypeVar("F")
  8
  9
 10class SegmentLazyQuadraticDivision(Generic[T, F]):
 11    """
 12    区間の総積取得・区間への作用適用クエリをそれぞれ時間計算量 :math:`O(\\sqrt{n})` で処理できるデータ構造です。
 13    定数倍が軽いのでそこまで遅くはないはずです。
 14    """
 15
 16    def __init__(
 17        self,
 18        n_or_a: Union[int, Iterable[T]],
 19        op: Callable[[T, T], T],
 20        mapping: Callable[[F, T], T],
 21        composition: Callable[[F, F], F],
 22        e: T,
 23        id: F,
 24    ) -> None:
 25        """
 26        引数は遅延セグ木のアレです。
 27
 28        :math:`O(n)` です。
 29        """
 30        if isinstance(n_or_a, int):
 31            self.n = n_or_a
 32            a = [e] * self.n
 33        else:
 34            if not isinstance(n_or_a, list):
 35                a = list(n_or_a)
 36            else:
 37                a = n_or_a
 38        self.n = len(a)
 39        self.op = op
 40        self.mapping = mapping
 41        self.composition = composition
 42        self.e = e
 43        self.id = id
 44        self.bucket_size = int(self.n**0.5) + 1
 45        self.bucket_cnt = (self.n + self.bucket_size - 1) // self.bucket_size
 46        self.data = [
 47            a[k * self.bucket_size : (k + 1) * self.bucket_size]
 48            for k in range(self.bucket_cnt)
 49        ]
 50        self.bucket_data = [reduce(self.op, v) for v in self.data]
 51        self.bucket_lazy = [id] * self.bucket_cnt
 52
 53    def apply(self, l: int, r: int, f: F) -> None:
 54        """区間 ``[l, r)`` に ``f`` を作用します。
 55
 56        :math:`O(\\sqrt{n})` です。
 57        """
 58        assert 0 <= l <= r <= self.n
 59
 60        def _change_data(k: int, l: int, r: int) -> None:
 61            self._propagate(k)
 62            d = self.data[k]
 63            for i in range(l, r):
 64                d[i] = self.mapping(f, d[i])
 65            self.bucket_data[k] = reduce(self.op, self.data[k])
 66
 67        k1 = l // self.bucket_size
 68        k2 = r // self.bucket_size
 69        l -= k1 * self.bucket_size
 70        r -= k2 * self.bucket_size
 71        if k1 == k2:
 72            if k1 < self.bucket_cnt:
 73                _change_data(k1, l, r)
 74        else:
 75            if k1 < self.bucket_cnt:
 76                if l == 0:
 77                    self.bucket_lazy[k1] = (
 78                        f
 79                        if self.bucket_lazy[k1] == self.id
 80                        else self.composition(f, self.bucket_lazy[k1])
 81                    )
 82                    self.bucket_data[k1] = self.mapping(f, self.bucket_data[k1])
 83                else:
 84                    _change_data(k1, l, len(self.data[k1]))
 85            for i in range(k1 + 1, k2):
 86                self.bucket_lazy[i] = (
 87                    f
 88                    if self.bucket_lazy[i] == self.id
 89                    else self.composition(f, self.bucket_lazy[i])
 90                )
 91                self.bucket_data[i] = self.mapping(f, self.bucket_data[i])
 92            if k2 < self.bucket_cnt:
 93                if r == len(self.data[k2]):
 94                    self.bucket_lazy[k2] = (
 95                        f
 96                        if self.bucket_lazy[k2] == self.id
 97                        else self.composition(f, self.bucket_lazy[k2])
 98                    )
 99                    self.bucket_data[k2] = self.mapping(f, self.bucket_data[k2])
100                else:
101                    _change_data(k2, 0, r)
102
103    def all_apply(self, f: F) -> None:
104        """区間 ``[0, n)`` に ``f`` を作用します。
105
106        :math:`O(\\sqrt{n})` です。
107        """
108        self.bucket_lazy = [
109            f if bl == self.id else self.composition(f, bl) for bl in self.bucket_lazy
110        ]
111
112    def _propagate(self, k: int) -> None:
113        if self.bucket_lazy[k] == self.id:
114            return
115        f = self.bucket_lazy[k]
116        dk = self.data[k]
117        for i, d in enumerate(dk):
118            dk[i] = self.mapping(f, d)
119        self.bucket_lazy[k] = self.id
120
121    def _all_propagatae(self) -> None:
122        for k in range(self.bucket_cnt):
123            self._propagate(k)
124        for i in range(self.bucket_cnt):
125            self.bucket_lazy[i] = self.id
126
127    def prod(self, l: int, r: int) -> T:
128        """区間 ``[l, r)`` の総積を返します。
129
130        :math:`O(\\sqrt{n})` です。
131
132        Args:
133          l (int): インデックスです。
134          r (int): インデックスです。
135        """
136        assert 0 <= l <= r <= self.n
137        if l == r:
138            return self.e
139        k1 = l // self.bucket_size
140        k2 = r // self.bucket_size
141        l -= k1 * self.bucket_size
142        r -= k2 * self.bucket_size
143        s = self.e
144        if k1 == k2:
145            s = reduce(self.op, self.data[k1][l:r])
146            if self.bucket_lazy[k1] != self.id:
147                s = self.mapping(self.bucket_lazy[k1], s)
148        else:
149            if l < len(self.data[k1]):
150                s = reduce(self.op, self.data[k1][l:])
151                if self.bucket_lazy[k1] != self.id:
152                    s = self.mapping(self.bucket_lazy[k1], s)
153            if k1 + 1 < k2:
154                s = reduce(self.op, self.bucket_data[k1 + 1 : k2], s)
155            if k2 < self.bucket_cnt and r > 0:
156                s_ = reduce(self.op, self.data[k2][:r])
157                if self.bucket_lazy[k2] != self.id:
158                    s_ = self.mapping(self.bucket_lazy[k2], s_)
159                s = self.op(s, s_)
160        return s
161
162    def all_prod(self) -> T:
163        """区間 ``[0, n)`` の総積を返します。
164
165        :math:`O(\\sqrt{n})` です。
166        """
167        return reduce(self.op, self.bucket_data)
168
169    def tolist(self) -> list[T]:
170        self._all_propagatae()
171        return list(chain(*self.data))
172
173    def __getitem__(self, k: int) -> T:
174        p = k // self.bucket_size
175        return (
176            self.data[p][k - p * self.bucket_size]
177            if self.bucket_lazy[p] == self.id
178            else self.mapping(
179                self.bucket_lazy[p], self.data[p][k - p * self.bucket_size]
180            )
181        )
182
183    def __setitem__(self, k, v):
184        p = k // self.bucket_size
185        self._propagate(p)
186        self.data[p][k - p * self.bucket_size] = v
187        self.bucket_data[p] = reduce(self.op, self.data[p])
188
189    def __str__(self):
190        return str(self.tolist())
191
192    def __repr__(self):
193        return f"SegmentLazyQuadraticDivision({self})"

仕様

class SegmentLazyQuadraticDivision(n_or_a: int | Iterable[T], op: Callable[[T, T], T], mapping: Callable[[F, T], T], composition: Callable[[F, F], F], e: T, id: F)[source]

Bases: Generic[T, F]

区間の総積取得・区間への作用適用クエリをそれぞれ時間計算量 O(n) で処理できるデータ構造です。 定数倍が軽いのでそこまで遅くはないはずです。

all_apply(f: F) None[source]

区間 [0, n)f を作用します。

O(n) です。

all_prod() T[source]

区間 [0, n) の総積を返します。

O(n) です。

apply(l: int, r: int, f: F) None[source]

区間 [l, r)f を作用します。

O(n) です。

prod(l: int, r: int) T[source]

区間 [l, r) の総積を返します。

O(n) です。

Parameters:
  • l (int) – インデックスです。

  • r (int) – インデックスです。

tolist() list[T][source]