diff --git a/README.md b/README.md index 93fc634..a3d7c01 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ ac-library-python is a Python port of [AtCoder Library (ACL)](https://atcoder.jp #### Data structure + [Fenwick Tree](https://github.com/atcoder/ac-library/blob/master/document_en/fenwicktree.md) ++ segtree #### Math @@ -32,7 +33,6 @@ ac-library-python is a Python port of [AtCoder Library (ACL)](https://atcoder.jp #### Data structure -+ segtree + lazysegtree + string diff --git a/README_ja.md b/README_ja.md index 89f4d32..dde9c9d 100644 --- a/README_ja.md +++ b/README_ja.md @@ -13,6 +13,7 @@ ac-library-pythonは、[AtCoder Library (ACL)](https://atcoder.jp/posts/517)のP #### データ構造 + [Fenwick Tree](https://github.com/atcoder/ac-library/blob/master/document_ja/fenwicktree.md) ++ segtree #### 数学 @@ -30,7 +31,6 @@ ac-library-pythonは、[AtCoder Library (ACL)](https://atcoder.jp/posts/517)のP #### データ構造 -+ segtree + lazysegtree + string diff --git a/atcoder/_bit.py b/atcoder/_bit.py new file mode 100644 index 0000000..9d1c939 --- /dev/null +++ b/atcoder/_bit.py @@ -0,0 +1,15 @@ +def _ceil_pow2(n: int) -> int: + x = 0 + while (1 << x) < n: + x += 1 + + return x + + +def _bsf(n: int) -> int: + x = 0 + while n % 2 == 0: + x += 1 + n //= 2 + + return x diff --git a/atcoder/segtree.py b/atcoder/segtree.py new file mode 100644 index 0000000..25c263e --- /dev/null +++ b/atcoder/segtree.py @@ -0,0 +1,119 @@ +import typing + +import atcoder._bit + + +class SegTree: + def __init__(self, + op: typing.Callable[[typing.Any, typing.Any], typing.Any], + e: typing.Any, + v: typing.Union[int, typing.List[typing.Any]]) -> None: + self._op = op + self._e = e + + if isinstance(v, int): + v = [e] * v + + self._n = len(v) + self._log = atcoder._bit._ceil_pow2(self._n) + self._size = 1 << self._log + self._d = [e] * (2 * self._size) + + for i in range(self._n): + self._d[self._size + i] = v[i] + for i in range(self._size - 1, 0, -1): + self._update(i) + + def set(self, p: int, x: typing.Any) -> None: + assert 0 <= p < self._n + + p += self._size + self._d[p] = x + for i in range(1, self._log + 1): + self._update(p >> i) + + def get(self, p: int) -> typing.Any: + assert 0 <= p < self._n + + return self._d[p + self._size] + + def prod(self, left: int, right: int) -> typing.Any: + assert 0 <= left <= right <= self._n + sml = self._e + smr = self._e + left += self._size + right += self._size + + while left < right: + if left & 1: + sml = self._op(sml, self._d[left]) + left += 1 + if right & 1: + right -= 1 + smr = self._op(self._d[right], smr) + left >>= 1 + right >>= 1 + + return self._op(sml, smr) + + def all_prod(self) -> typing.Any: + return self._d[1] + + def max_right(self, left: int, + f: typing.Callable[[typing.Any], bool]) -> int: + assert 0 <= left <= self._n + assert f(self._e) + + if left == self._n: + return self._n + + left += self._size + sm = self._e + + first = True + while first or (left & -left) != left: + first = False + while left % 2 == 0: + left >>= 1 + if not f(self._op(sm, self._d[left])): + while left < self._size: + left *= 2 + if f(self._op(sm, self._d[left])): + sm = self._op(sm, self._d[left]) + left += 1 + return left - self._size + sm = self._op(sm, self._d[left]) + left += 1 + + return self._n + + def min_left(self, right: int, + f: typing.Callable[[typing.Any], bool]) -> int: + assert 0 <= right <= self._n + assert f(self._e) + + if right == 0: + return 0 + + right += self._size + sm = self._e + + first = True + while first or (right & -right) != right: + first = False + right -= 1 + while right > 1 and right % 2: + right >>= 1 + if not f(self._op(self._d[right], sm)): + while right < self._size: + right = 2 * right + 1 + if f(self._op(self._d[right], sm)): + sm = self._op(self._d[right], sm) + right -= 1 + return right + 1 - self._size + sm = self._op(self._d[right], sm) + + return 0 + + def _update(self, k: int) -> None: + self._d[k] = self._op(self._d[2 * k], self._d[2 * k + 1]) diff --git a/example/segtree_practice.py b/example/segtree_practice.py new file mode 100644 index 0000000..b0a53d7 --- /dev/null +++ b/example/segtree_practice.py @@ -0,0 +1,25 @@ +# https://atcoder.jp/contests/practice2/tasks/practice2_j + +import sys + +from atcoder.segtree import SegTree + + +def main() -> None: + n, q = map(int, sys.stdin.readline().split()) + a = list(map(int, sys.stdin.readline().split())) + + segtree = SegTree(max, -1, a) + + for _ in range(q): + t, x, y = map(int, sys.stdin.readline().split()) + if t == 1: + segtree.set(x - 1, y) + elif t == 2: + print(segtree.prod(x - 1, y)) + else: + print(segtree.max_right(x - 1, lambda v: v < y) + 1) + + +if __name__ == '__main__': + main() diff --git a/example/segtree_practice_reversed.py b/example/segtree_practice_reversed.py new file mode 100644 index 0000000..3e26259 --- /dev/null +++ b/example/segtree_practice_reversed.py @@ -0,0 +1,25 @@ +# https://atcoder.jp/contests/practice2/tasks/practice2_j + +import sys + +from atcoder.segtree import SegTree + + +def main() -> None: + n, q = map(int, sys.stdin.readline().split()) + a = list(map(int, sys.stdin.readline().split()))[::-1] + + segtree = SegTree(max, -1, a) + + for _ in range(q): + t, x, y = map(int, sys.stdin.readline().split()) + if t == 1: + segtree.set(n - x, y) + elif t == 2: + print(segtree.prod(n - y, n - x + 1)) + else: + print(n - segtree.min_left(n - x + 1, lambda v: v < y) + 1) + + +if __name__ == '__main__': + main()