From c2ac2df6013b27879fc80b6fcb3ee443435a1052 Mon Sep 17 00:00:00 2001 From: Naoto Mizuno Date: Sat, 19 Sep 2020 09:25:19 +0900 Subject: [PATCH] Port math --- README.md | 5 +- README_ja.md | 5 +- atcoder/_math.py | 106 ++++++++++++++++++++++++++++++++++ atcoder/math.py | 93 +++++++++++++++++++++++++++++ example/floor_sum_practice.py | 16 +++++ 5 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 atcoder/_math.py create mode 100644 atcoder/math.py create mode 100644 example/floor_sum_practice.py diff --git a/README.md b/README.md index 33caeaa..2abe4d9 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,10 @@ 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) +#### Math + ++ math + #### Graph + [Disjoint Set Union (DSU)](https://github.com/atcoder/ac-library/blob/master/document_en/dsu.md) @@ -30,7 +34,6 @@ ac-library-python is a Python port of [AtCoder Library (ACL)](https://atcoder.jp #### Math -+ math + convolution + modint diff --git a/README_ja.md b/README_ja.md index ba472c1..0c29fc1 100644 --- a/README_ja.md +++ b/README_ja.md @@ -14,6 +14,10 @@ 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) +#### 数学 + ++ math + #### グラフ + [Disjoint Set Union (DSU)](https://github.com/atcoder/ac-library/blob/master/document_ja/dsu.md) @@ -28,7 +32,6 @@ ac-library-pythonは、[AtCoder Library (ACL)](https://atcoder.jp/posts/517)のP #### 数学 -+ math + convolution + modint diff --git a/atcoder/_math.py b/atcoder/_math.py new file mode 100644 index 0000000..7a3605c --- /dev/null +++ b/atcoder/_math.py @@ -0,0 +1,106 @@ +import typing + + +def _is_prime(n: int) -> bool: + ''' + Reference: + M. Forisek and J. Jancina, + Fast Primality Testing for Integers That Fit into a Machine Word + ''' + + if n <= 1: + return False + if n == 2 or n == 7 or n == 61: + return True + if n % 2 == 0: + return False + + d = n - 1 + while d % 2 == 0: + d //= 2 + + for a in (2, 7, 61): + t = d + y = pow(a, t, n) + while t != n - 1 and y != 1 and y != n - 1: + y = y * y % n + t <<= 1 + if y != n - 1 and t % 2 == 0: + return False + return True + + +def _inv_gcd(a: int, b: int) -> typing.Tuple[int, int]: + a %= b + if a == 0: + return (b, 0) + + # Contracts: + # [1] s - m0 * a = 0 (mod b) + # [2] t - m1 * a = 0 (mod b) + # [3] s * |m1| + t * |m0| <= b + s = b + t = a + m0 = 0 + m1 = 1 + + while t: + u = s // t + s -= t * u + m0 -= m1 * u # |m1 * u| <= |m1| * s <= b + + # [3]: + # (s - t * u) * |m1| + t * |m0 - m1 * u| + # <= s * |m1| - t * u * |m1| + t * (|m0| + |m1| * u) + # = s * |m1| + t * |m0| <= b + + s, t = t, s + m0, m1 = m1, m0 + + # by [3]: |m0| <= b/g + # by g != b: |m0| < b/g + if m0 < 0: + m0 += b // s + + return (s, m0) + + +def _primitive_root(m: int) -> int: + if m == 2: + return 1 + if m == 167772161: + return 3 + if m == 469762049: + return 3 + if m == 754974721: + return 11 + if m == 998244353: + return 3 + + divs = [2] + [0] * 19 + cnt = 1 + x = (m - 1) // 2 + while x % 2 == 0: + x //= 2 + + i = 3 + while i * i <= x: + if x % i == 0: + divs[cnt] = i + cnt += 1 + while x % i == 0: + x //= i + i += 2 + + if x > 1: + divs[cnt] = x + cnt += 1 + + g = 2 + while True: + for i in range(cnt): + if pow(g, (m - 1) // divs[i], m) == 1: + break + else: + return g + g += 1 diff --git a/atcoder/math.py b/atcoder/math.py new file mode 100644 index 0000000..7f2c689 --- /dev/null +++ b/atcoder/math.py @@ -0,0 +1,93 @@ +import typing + +import atcoder._math + + +def inv_mod(x: int, m: int) -> int: + assert 1 <= m + + z = atcoder._math._inv_gcd(x, m) + + assert z[0] == 1 + + return z[1] + + +def crt(r: typing.List[int], m: typing.List[int]) -> typing.Tuple[int, int]: + assert len(r) == len(m) + + n = len(r) + + # Contracts: 0 <= r0 < m0 + r0 = 0 + m0 = 1 + for i in range(n): + assert 1 <= m[i] + r1 = r[i] % m[i] + m1 = m[i] + if m0 < m1: + r0, r1 = r1, r0 + m0, m1 = m1, m0 + if m0 % m1 == 0: + if r0 % m1 != r1: + return (0, 0) + continue + + # assume: m0 > m1, lcm(m0, m1) >= 2 * max(m0, m1) + + ''' + (r0, m0), (r1, m1) -> (r2, m2 = lcm(m0, m1)); + r2 % m0 = r0 + r2 % m1 = r1 + -> (r0 + x*m0) % m1 = r1 + -> x*u0*g % (u1*g) = (r1 - r0) (u0*g = m0, u1*g = m1) + -> x = (r1 - r0) / g * inv(u0) (mod u1) + ''' + + # im = inv(u0) (mod u1) (0 <= im < u1) + g, im = atcoder._math._inv_gcd(m0, m1) + + u1 = m1 // g + # |r1 - r0| < (m0 + m1) <= lcm(m0, m1) + if (r1 - r0) % g: + return (0, 0) + + # u1 * u1 <= m1 * m1 / g / g <= m0 * m1 / g = lcm(m0, m1) + x = (r1 - r0) // g % u1 * im % u1 + + ''' + |r0| + |m0 * x| + < m0 + m0 * (u1 - 1) + = m0 + m0 * m1 / g - m0 + = lcm(m0, m1) + ''' + + r0 += x * m0 + m0 *= u1 # -> lcm(m0, m1) + if r0 < 0: + r0 += m0 + + return (r0, m0) + + +def floor_sum(n: int, m: int, a: int, b: int) -> int: + ans = 0 + + if a >= m: + ans += (n - 1) * n * (a // m) // 2 + a %= m + + if b >= m: + ans += n * (b // m) + b %= m + + y_max = (a * n + b) // m + x_max = y_max * m - b + + if y_max == 0: + return ans + + ans += (n - (x_max + a - 1) // a) * y_max + ans += floor_sum(y_max, a, m, (a - x_max % a) % a) + + return ans diff --git a/example/floor_sum_practice.py b/example/floor_sum_practice.py new file mode 100644 index 0000000..00dfc56 --- /dev/null +++ b/example/floor_sum_practice.py @@ -0,0 +1,16 @@ +# https://atcoder.jp/contests/practice2/tasks/practice2_c + +from atcoder.math import floor_sum + + +def main() -> None: + import sys + + t = int(sys.stdin.readline()) + for _ in range(t): + n, m, a, b = map(int, sys.stdin.readline().split()) + print(floor_sum(n, m, a, b)) + + +if __name__ == '__main__': + main()