1from typing import Union, Callable, TypeVar, Generic, Iterable
2
3T = TypeVar("T")
4F = TypeVar("F")
5
6
[docs]
7class DualCommutativeSegmentTree(Generic[T, F]):
8 """作用が可換な双対セグ木です。"""
9
10 def __init__(
11 self,
12 n_or_a: Union[int, Iterable[T]],
13 mapping: Callable[[F, T], T],
14 composition: Callable[[F, F], F],
15 e: T,
16 id: F,
17 ) -> None:
18 self.mapping: Callable[[F, T], T] = mapping
19 self.composition: Callable[[F, F], F] = composition
20 self.e: T = e
21 self.id: F = id
22 self.data: list[T] = [e] * n_or_a if isinstance(n_or_a, int) else list(n_or_a)
23 self.n: int = len(self.data)
24 self.log: int = (self.n - 1).bit_length()
25 self.size: int = 1 << self.log
26 self.lazy: list[F] = [id] * self.size
27
28 def _all_apply(self, k: int, f: F) -> None:
29 if k < self.size:
30 self.lazy[k] = self.composition(f, self.lazy[k])
31 return
32 k -= self.size
33 if k < self.n:
34 self.data[k] = self.mapping(f, self.data[k])
35
36 def _propagate(self, k: int) -> None:
37 self._all_apply(k << 1, self.lazy[k])
38 self._all_apply(k << 1 | 1, self.lazy[k])
39 self.lazy[k] = self.id
40
[docs]
41 def apply_point(self, k: int, f: F) -> None:
42 """``k`` 番目の値に ``f`` を作用させます。
43 :math:`O(1)` です。
44
45 Args:
46 k (int): _description_
47 f (F): _description_
48 """
49 assert (
50 0 <= k < self.n
51 ), f"IndexError: {self.__class__.__name__}.apply_point({k}, {f}), n={self.n}"
52 self.data[k] = self.mapping(f, self.data[k])
53
[docs]
54 def apply(self, l: int, r: int, f: F) -> None:
55 """区間 ``[l, r)`` に ``f`` を作用させます。
56 :math:`O(\\log{n})` です。
57
58 Args:
59 l (int): _description_
60 r (int): _description_
61 f (F): _description_
62 """
63 assert (
64 0 <= l <= r <= self.n
65 ), f"IndexError: {self.__class__.__name__}.apply({l}, {r}, {f}), n={self.n}"
66 if l == r:
67 return
68 if f == self.id:
69 return
70 l += self.size
71 r += self.size
72 data, lazy = self.data, self.lazy
73 if (l - self.size) & 1:
74 data[l - self.size] = self.mapping(f, data[l - self.size])
75 l += 1
76 if (r - self.size) & 1:
77 data[(r - self.size) ^ 1] = self.mapping(f, data[(r - self.size) ^ 1])
78 r ^= 1
79 l >>= 1
80 r >>= 1
81 while l < r:
82 if l & 1:
83 lazy[l] = self.composition(f, lazy[l])
84 l += 1
85 if r & 1:
86 r ^= 1
87 lazy[r] = self.composition(f, lazy[r])
88 l >>= 1
89 r >>= 1
90
[docs]
91 def all_apply(self, f: F) -> None:
92 """区間 ``[0, n)`` に ``f`` を作用させます。
93 :math:`O(1)` です。
94 """
95 self.lazy[1] = self.composition(f, self.lazy[1])
96
[docs]
97 def all_propagate(self) -> None:
98 for i in range(self.size):
99 self._propagate(i)
100
[docs]
101 def tolist(self) -> list[T]:
102 self.all_propagate()
103 return self.data[:]
104
[docs]
105 def __getitem__(self, k: int) -> T:
106 """``k`` 番目の値を取得します。
107 :math:`O(\\log{n})` です。
108 """
109 assert (
110 -self.n <= k < self.n
111 ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
112 if k < 0:
113 k += self.n
114 lazy = self.lazy
115 k += self.size
116 lazy_f = self.id
117 for i in range(self.log, 0, -1):
118 lazy_f = self.composition(lazy_f, lazy[k >> i])
119 return self.mapping(lazy_f, self.data[k - self.size])
120
121 def __setitem__(self, k: int, v: T) -> None:
122 assert (
123 -self.n <= k < self.n
124 ), f"IndexError: {self.__class__.__name__}[{k}] = {v}, n={self.n}"
125 if k < 0:
126 k += self.n
127 k += self.size
128 for i in range(self.log, 0, -1):
129 self._propagate(k >> i)
130 self.data[k - self.size] = v
131
132 def __str__(self):
133 return str([self[i] for i in range(self.n)])
134
135 def __repr__(self):
136 return f"{self.__class__.__name__}({self})"