Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ ac-library-python is a Python port of [AtCoder Library (ACL)](https://atcoder.jp

+ [Fenwick Tree](https://github.com/atcoder/ac-library/blob/master/document_en/fenwicktree.md)
+ segtree
+ lazysegtree

#### Math

Expand All @@ -34,7 +35,6 @@ ac-library-python is a Python port of [AtCoder Library (ACL)](https://atcoder.jp

#### Data structure

+ lazysegtree
+ string

#### Math
Expand Down
2 changes: 1 addition & 1 deletion README_ja.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,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
+ lazysegtree

#### 数学

Expand All @@ -32,7 +33,6 @@ ac-library-pythonは、[AtCoder Library (ACL)](https://atcoder.jp/posts/517)のP

#### データ構造

+ lazysegtree
+ string

#### 数学
Expand Down
204 changes: 204 additions & 0 deletions atcoder/lazysegtree.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,204 @@
import typing

import atcoder._bit


class LazySegTree:
def __init__(
self,
op: typing.Callable[[typing.Any, typing.Any], typing.Any],
e: typing.Any,
mapping: typing.Callable[[typing.Any, typing.Any], typing.Any],
composition: typing.Callable[[typing.Any, typing.Any], typing.Any],
id_: typing.Any,
v: typing.Union[int, typing.List[typing.Any]]) -> None:
self._op = op
self._e = e
self._mapping = mapping
self._composition = composition
self._id = id_

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)
self._lz = [self._id] * 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
for i in range(self._log, 0, -1):
self._push(p >> i)
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

p += self._size
for i in range(self._log, 0, -1):
self._push(p >> i)
return self._d[p]

def prod(self, left: int, right: int) -> typing.Any:
assert 0 <= left <= right <= self._n

if left == right:
return self._e

left += self._size
right += self._size

for i in range(self._log, 0, -1):
if ((left >> i) << i) != left:
self._push(left >> i)
if ((right >> i) << i) != right:
self._push(right >> i)

sml = self._e
smr = self._e
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 apply(self, left: int, right: typing.Optional[int] = None,
f: typing.Optional[typing.Any] = None):
assert f is not None

if right is None:
p = left
assert 0 <= left < self._n

p += self._size
for i in range(self._log, 0, -1):
self._push(p >> i)
self._d[p] = self._mapping(f, self._d[p])
for i in range(1, self._log + 1):
self._update(p >> i)
else:
assert 0 <= left <= right <= self._n
if left == right:
return

left += self._size
right += self._size

for i in range(self._log, 0, -1):
if ((left >> i) << i) != left:
self._push(left >> i)
if ((right >> i) << i) != right:
self._push((right - 1) >> i)

l2 = left
r2 = right
while left < right:
if left & 1:
self._all_apply(left, f)
left += 1
if right & 1:
right -= 1
self._all_apply(right, f)
left >>= 1
right >>= 1
left = l2
right = r2

for i in range(1, self._log + 1):
if ((left >> i) << i) != left:
self._update(left >> i)
if ((right >> i) << i) != right:
self._update((right - 1) >> i)

def max_right(
self, left: int, g: typing.Callable[[typing.Any], bool]) -> int:
assert 0 <= left <= self._n
assert g(self._e)

if left == self._n:
return self._n

left += self._size
for i in range(self._log, 0, -1):
self._push(left >> i)

sm = self._e
first = True
while first or (left & -left) != left:
first = False
while left % 2 == 0:
left >>= 1
if not g(self._op(sm, self._d[left])):
while left < self._size:
self._push(left)
left *= 2
if g(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, g: typing.Any) -> int:
assert 0 <= right <= self._n
assert g(self._e)

if right == 0:
return 0

right += self._size
for i in range(self._log, 0, -1):
self._push((right - 1) >> i)

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 g(self._op(self._d[right], sm)):
while right < self._size:
self._push(right)
right = 2 * right + 1
if g(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])

def _all_apply(self, k: int, f: typing.Any) -> None:
self._d[k] = self._mapping(f, self._d[k])
if k < self._size:
self._lz[k] = self._composition(f, self._lz[k])

def _push(self, k: int) -> None:
self._all_apply(2 * k, self._lz[k])
self._all_apply(2 * k + 1, self._lz[k])
self._lz[k] = self._id
79 changes: 46 additions & 33 deletions atcoder/modint.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from __future__ import annotations
import copy
import typing

import atcoder._math
Expand Down Expand Up @@ -37,27 +36,36 @@ def val(self) -> int:
return self._v

def __iadd__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
self._v += rhs._v
if isinstance(rhs, Modint):
self._v += rhs._v
else:
self._v += rhs
if self._v >= self._mod:
self._v -= self._mod
return self

def __isub__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
self._v -= rhs._v
if isinstance(rhs, Modint):
self._v -= rhs._v
else:
self._v -= rhs
if self._v < 0:
self._v += self._mod
return self

def __imul__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
self._v = self._v * rhs._v % self._mod
if isinstance(rhs, Modint):
self._v = self._v * rhs._v % self._mod
else:
self._v = self._v * rhs % self._mod
return self

def __ifloordiv__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
self *= rhs.inv()
if isinstance(rhs, Modint):
inv = rhs.inv()._v
else:
inv = atcoder._math._inv_gcd(rhs, self._mod)[1]
self._v = self._v * inv % self._mod
return self

def __pos__(self) -> Modint:
Expand All @@ -79,42 +87,47 @@ def inv(self) -> Modint:
return Modint(eg[1])

def __add__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
result = copy.deepcopy(self)
result += rhs
return result
if isinstance(rhs, Modint):
result = self._v + rhs._v
if result >= self._mod:
result -= self._mod
return raw(result)
else:
return Modint(self._v + rhs)

def __sub__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
result = copy.deepcopy(self)
result -= rhs
return result
if isinstance(rhs, Modint):
result = self._v - rhs._v
if result < 0:
result += self._mod
return raw(result)
else:
return Modint(self._v - rhs)

def __mul__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
result = copy.deepcopy(self)
result *= rhs
return result
if isinstance(rhs, Modint):
return Modint(self._v * rhs._v)
else:
return Modint(self._v * rhs)

def __floordiv__(self, rhs: typing.Union[Modint, int]) -> Modint:
rhs = self._asmodint(rhs)
result = copy.deepcopy(self)
result //= rhs
return result
if isinstance(rhs, Modint):
inv = rhs.inv()._v
else:
inv = atcoder._math._inv_gcd(rhs, self._mod)[1]
return Modint(self._v * inv)

def __eq__(self, rhs: typing.Union[Modint, int]) -> bool:
rhs = self._asmodint(rhs)
return self._v == rhs._v
if isinstance(rhs, Modint):
return self._v == rhs._v
else:
return self._v == rhs

def __ne__(self, rhs: typing.Union[Modint, int]) -> bool:
rhs = self._asmodint(rhs)
return self._v != rhs._v

def _asmodint(self, rhs: typing.Union[Modint, int]) -> Modint:
if isinstance(rhs, Modint):
return rhs
return self._v != rhs._v
else:
return Modint(rhs)
return self._v != rhs


def raw(v: int) -> Modint:
Expand Down
Loading