1from typing import Union, Callable, TypeVar, Generic, Iterable
2
3T = TypeVar("T")
4F = TypeVar("F")
5
6
[docs]
7class DualSegmentTree(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 assert (
43 0 <= k < self.n
44 ), f"IndexError: {self.__class__.__name__}.apply_point({k}, {f}), n={self.n}"
45 k += self.size
46 for i in range(self.log, 0, -1):
47 self._propagate(k >> i)
48 self.data[k - self.size] = self.mapping(f, self.data[k - self.size])
49
[docs]
50 def apply(self, l: int, r: int, f: F) -> None:
51 assert (
52 0 <= l <= r <= self.n
53 ), f"IndexError: {self.__class__.__name__}.apply({l}, {r}, {f}), n={self.n}"
54 if l == r:
55 return
56 if f == self.id:
57 return
58 l += self.size
59 r += self.size
60 data, lazy = self.data, self.lazy
61 for i in range(self.log, 0, -1):
62 if l >> i << i != l and lazy[l >> i] != self.id:
63 self._propagate(l >> i)
64 if r >> i << i != r and lazy[(r - 1) >> i] != self.id:
65 self._propagate((r - 1) >> i)
66 if (l - self.size) & 1:
67 data[l - self.size] = self.mapping(f, data[l - self.size])
68 l += 1
69 if (r - self.size) & 1:
70 data[(r - self.size) ^ 1] = self.mapping(f, data[(r - self.size) ^ 1])
71 r ^= 1
72 l >>= 1
73 r >>= 1
74 while l < r:
75 if l & 1:
76 lazy[l] = self.composition(f, lazy[l])
77 l += 1
78 if r & 1:
79 r ^= 1
80 lazy[r] = self.composition(f, lazy[r])
81 l >>= 1
82 r >>= 1
83
[docs]
84 def all_apply(self, f: F) -> None:
85 self.lazy[1] = self.composition(f, self.lazy[1])
86
[docs]
87 def all_propagate(self) -> None:
88 for i in range(self.size):
89 self._propagate(i)
90
[docs]
91 def tolist(self) -> list[T]:
92 self.all_propagate()
93 return self.data[:]
94
95 def __getitem__(self, k: int) -> T:
96 assert (
97 -self.n <= k < self.n
98 ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
99 if k < 0:
100 k += self.n
101 k += self.size
102 for i in range(self.log, 0, -1):
103 self._propagate(k >> i)
104 return self.data[k - self.size]
105
106 def __setitem__(self, k: int, v: T) -> None:
107 assert (
108 -self.n <= k < self.n
109 ), f"IndexError: {self.__class__.__name__}[{k}] = {v}, n={self.n}"
110 if k < 0:
111 k += self.n
112 k += self.size
113 for i in range(self.log, 0, -1):
114 self._propagate(k >> i)
115 self.data[k - self.size] = v
116
117 def __str__(self):
118 return str([self[i] for i in range(self.n)])
119
120 def __repr__(self):
121 return f"{self.__class__.__name__}({self})"