1from titan_pylib.data_structures.segment_tree.segment_tree_interface import (
2 SegmentTreeInterface,
3)
4from typing import Generic, TypeVar, Callable
5
6T = TypeVar("T")
7
8
[docs]
9class DynamicSegmentTree(SegmentTreeInterface, Generic[T]):
10 """動的セグ木です。"""
11
12 def __init__(self, u: int, op: Callable[[T, T], T], e: T):
13 self._op = op
14 self._e = e
15 self._u = u
16 self._log = (self._u - 1).bit_length()
17 self._size = 1 << self._log
18 self._data: dict[int, T] = {}
19
[docs]
20 def set(self, k: int, v: T) -> None:
21 assert (
22 -self._u <= k < self._u
23 ), f"IndexError: {self.__class__.__name__}.set({k}: int, {v}: T), n={self._u}"
24 if k < 0:
25 k += self._u
26 k += self._size
27 self._data[k] = v
28 e = self._e
29 for _ in range(self._log):
30 k >>= 1
31 self._data[k] = self._op(
32 self._data.get(k << 1, e), self._data.get(k << 1 | 1, e)
33 )
34
[docs]
35 def get(self, k: int) -> T:
36 assert (
37 -self._u <= k < self._u
38 ), f"IndexError: {self.__class__.__name__}.get({k}: int), n={self._u}"
39 if k < 0:
40 k += self._u
41 return self._data.get(k + self._size, self._e)
42
[docs]
43 def prod(self, l: int, r: int) -> T:
44 assert (
45 0 <= l <= r <= self._u
46 ), f"IndexError: {self.__class__.__name__}.prod({l}: int, {r}: int)"
47 l += self._size
48 r += self._size
49 e = self._e
50 lres = e
51 rres = e
52 while l < r:
53 if l & 1:
54 lres = self._op(lres, self._data.get(l, e))
55 l += 1
56 if r & 1:
57 rres = self._op(self._data.get(r ^ 1, e), rres)
58 l >>= 1
59 r >>= 1
60 return self._op(lres, rres)
61
[docs]
62 def all_prod(self) -> T:
63 return self._data[1]
64
[docs]
65 def max_right(self, l: int, f: Callable[[T], bool]) -> int:
66 """Find the largest index R s.t. f([l, R)) == True. / O(logU)"""
67 assert (
68 0 <= l <= self._u
69 ), f"IndexError: {self.__class__.__name__}.max_right({l}, f) index out of range"
70 assert f(
71 self._e
72 ), f"{self.__class__.__name__}.max_right({l}, f), f({self._e}) must be true."
73 if l == self._u:
74 return self._u
75 l += self._size
76 e = self._e
77 s = e
78 while True:
79 while l & 1 == 0:
80 l >>= 1
81 if not f(self._op(s, self._data.get(l, e))):
82 while l < self._size:
83 l <<= 1
84 if f(self._op(s, self._data.get(l, e))):
85 s = self._op(s, self._data.get(l, e))
86 l |= 1
87 return l - self._size
88 s = self._op(s, self._data.get(l, e))
89 l += 1
90 if l & -l == l:
91 break
92 return self._u
93
[docs]
94 def min_left(self, r: int, f: Callable[[T], bool]) -> int:
95 """Find the smallest index L s.t. f([L, r)) == True. / O(logU)"""
96 assert (
97 0 <= r <= self._u
98 ), f"IndexError: {self.__class__.__name__}.min_left({r}, f) index out of range"
99 assert f(
100 self._e
101 ), f"{self.__class__.__name__}.min_left({r}, f), f({self._e}) must be true."
102 if r == 0:
103 return 0
104 r += self._size
105 e = self._e
106 s = e
107 while True:
108 r -= 1
109 while r > 1 and r & 1:
110 r >>= 1
111 if not f(self._op(self._data.get(r, e), s)):
112 while r < self._size:
113 r = r << 1 | 1
114 if f(self._op(self._data.get(r, e), s)):
115 s = self._op(self._data.get(r, e), s)
116 r ^= 1
117 return r + 1 - self._size
118 s = self._op(self._data.get(r, e), s)
119 if r & -r == r:
120 break
121 return 0
122
[docs]
123 def tolist(self) -> list[T]:
124 return [self.get(i) for i in range(self._u)]
125
126 def __getitem__(self, k: int) -> T:
127 assert (
128 -self._u <= k < self._u
129 ), f"IndexError: {self.__class__.__name__}[{k}]: int), n={self._u}"
130 return self.get(k)
131
132 def __setitem__(self, k: int, v: T) -> None:
133 assert (
134 -self._u <= k < self._u
135 ), f"IndexError: {self.__class__.__name__}.__setitem__{k}: int, {v}: T), n={self._u}"
136 self.set(k, v)
137
138 def __str__(self) -> str:
139 return str(self.tolist())
140
141 def __repr__(self) -> str:
142 return f"{self.__class__.__name__}({self})"