1from typing import Union, Callable, TypeVar, Generic, Iterable
2
3T = TypeVar("T")
4F = TypeVar("F")
5
6
[docs]
7class LazySegmentTreeRange(Generic[T, F]):
8 """遅延セグ木です。"""
9
10 def __init__(
11 self,
12 n_or_a: Union[int, Iterable[T]],
13 op: Callable[[T, T, int], T],
14 mapping: Callable[[F, T, int], T],
15 composition: Callable[[F, F, int], F],
16 e: T,
17 id: F,
18 ) -> None:
19 self.op: Callable[[T, T, int], T] = op
20 self.mapping: Callable[[F, T, int], T] = mapping
21 self.composition: Callable[[F, F, int], F] = composition
22 self.e: T = e
23 self.id: F = id
24 if isinstance(n_or_a, int):
25 self.n = n_or_a
26 self.log = (self.n - 1).bit_length()
27 self.size = 1 << self.log
28 size_data: list[int] = [1] * (self.size << 1)
29 for i in range(self.size - 1, 0, -1):
30 size_data[i] = size_data[i << 1] + size_data[i << 1 | 1]
31 self.size_data = size_data
32 self.data = [e] * (self.size << 1)
33 else:
34 a = list(n_or_a)
35 self.n = len(a)
36 self.log = (self.n - 1).bit_length()
37 self.size = 1 << self.log
38 data = [e] * (self.size << 1)
39 data[self.size : self.size + self.n] = a
40 size_data: list[int] = [1] * (self.size << 1)
41 for i in range(self.size - 1, 0, -1):
42 size_data[i] = size_data[i << 1] + size_data[i << 1 | 1]
43 data[i] = op(data[i << 1], data[i << 1 | 1], size_data[i])
44 self.size_data = size_data
45 self.data = data
46 self.lazy = [id] * self.size
47
48 def _update(self, k: int) -> None:
49 self.data[k] = self.op(
50 self.data[k << 1], self.data[k << 1 | 1], self.size_data[k]
51 )
52
53 def _all_apply(self, k: int, f: F) -> None:
54 self.data[k] = self.mapping(f, self.data[k], self.size_data[k])
55 if k >= self.size:
56 return
57 self.lazy[k] = self.composition(f, self.lazy[k], self.size_data[k])
58
59 def _propagate(self, k: int) -> None:
60 self._all_apply(k << 1, self.lazy[k])
61 self._all_apply(k << 1 | 1, self.lazy[k])
62 self.lazy[k] = self.id
63
[docs]
64 def apply_point(self, k: int, f: F) -> None:
65 k += self.size
66 for i in range(self.log, 0, -1):
67 self._propagate(k >> i)
68 self.data[k] = self.mapping(f, self.data[k], self.size_data[k])
69 for i in range(1, self.log + 1):
70 self._update(k >> i)
71
[docs]
72 def apply(self, l: int, r: int, f: F) -> None:
73 assert (
74 0 <= l <= r <= self.n
75 ), f"IndexError: {self.__class__.__name__}.apply({l}, {r}, {f}), n={self.n}"
76 if l == r:
77 return
78 if f == self.id:
79 return
80 l += self.size
81 r += self.size
82 lazy = self.lazy
83 for i in range(self.log, 0, -1):
84 if l >> i << i != l and lazy[l >> i] != self.id:
85 self._propagate(l >> i)
86 if r >> i << i != r and lazy[(r - 1) >> i] != self.id:
87 self._propagate((r - 1) >> i)
88 l2, r2 = l, r
89 while l < r:
90 if l & 1:
91 self._all_apply(l, f)
92 l += 1
93 if r & 1:
94 self._all_apply(r ^ 1, f)
95 l >>= 1
96 r >>= 1
97 ll, rr = l2, r2 - 1
98 for i in range(1, self.log + 1):
99 ll >>= 1
100 rr >>= 1
101 if ll << i != l2:
102 self._update(ll)
103 if r2 >> i << i != r2:
104 self._update(rr)
105
[docs]
106 def all_apply(self, f: F) -> None:
107 self.lazy[1] = self.composition(f, self.lazy[1], self.size_data[1])
108
[docs]
109 def prod(self, l: int, r: int) -> T:
110 assert (
111 0 <= l <= r <= self.n
112 ), f"IndexError: {self.__class__.__name__}.prod({l}, {r}), n={self.n}"
113 if l == r:
114 return self.e
115 l += self.size
116 r += self.size
117 lazy = self.lazy
118 for i in range(self.log, 0, -1):
119 ll, rr = l >> i, r >> i
120 if ll << i != l and lazy[ll] != self.id:
121 self._propagate(ll)
122 if rr << i != r and lazy[rr] != self.id:
123 self._propagate(rr)
124 lres = self.e
125 rres = self.e
126 lsize = 1
127 rsize = 1
128 while l < r:
129 if l & 1:
130 lsize += self.size_data[l]
131 lres = self.op(lres, self.data[l], lsize)
132 l += 1
133 if r & 1:
134 rsize += self.size_data[r ^ 1]
135 rres = self.op(self.data[r ^ 1], rres, rsize)
136 l >>= 1
137 r >>= 1
138 return self.op(lres, rres, lsize + rsize)
139
[docs]
140 def all_prod(self) -> T:
141 return self.data[1]
142
[docs]
143 def all_propagate(self) -> None:
144 for i in range(self.size):
145 self._propagate(i)
146
[docs]
147 def tolist(self) -> list[T]:
148 self.all_propagate()
149 return self.data[self.size : self.size + self.n]
150
151 def __getitem__(self, k: int) -> T:
152 assert (
153 -self.n <= k < self.n
154 ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
155 if k < 0:
156 k += self.n
157 k += self.size
158 for i in range(self.log, 0, -1):
159 self._propagate(k >> i)
160 return self.data[k]
161
162 def __setitem__(self, k: int, v: T):
163 assert (
164 -self.n <= k < self.n
165 ), f"IndexError: {self.__class__.__name__}[{k}] = {v}, n={self.n}"
166 if k < 0:
167 k += self.n
168 k += self.size
169 for i in range(self.log, 0, -1):
170 self._propagate(k >> i)
171 self.data[k] = v
172 for i in range(1, self.log + 1):
173 self._update(k >> i)
174
175 def __str__(self) -> str:
176 return "[" + ", ".join(map(str, (self[i] for i in range(self.n)))) + "]"
177
178 def __repr__(self):
179 return f"{self.__class__.__name__}({self})"