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