-
Notifications
You must be signed in to change notification settings - Fork 312
/
chinese_remainder.py
39 lines (31 loc) · 1.05 KB
/
chinese_remainder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import operator as op
from functools import reduce
def gcd(x, y):
"""greatest common divisor of x and y"""
while y:
x, y = y, x % y
return x
def chinese_remainder(a, p):
"""returns x s.t. x = a[i] (mod p[i]) where p[i] is prime for all i"""
prod = reduce(op.mul, p, 1)
x = [prod // pi for pi in p]
return sum(a[i] * pow(x[i], p[i] - 2, p[i]) * x[i] for i in range(len(a))) % prod
def extended_gcd(a, b):
"""returns gcd(a, b), s, r s.t. a * s + b * r == gcd(a, b)"""
s, old_s = 0, 1
r, old_r = b, a
while r:
q = old_r // r
old_r, r = r, old_r - q * r
old_s, s = s, old_s - q * s
return old_r, old_s, (old_r - old_s * a) // b if b else 0
def composite_crt(b, m):
"""returns x s.t. x = b[i] (mod m[i]) for all i"""
x, m_prod = 0, 1
for bi, mi in zip(b, m):
g, s, _ = extended_gcd(m_prod, mi)
if ((bi - x) % mi) % g:
return None
x += m_prod * (s * ((bi - x) % mi) // g)
m_prod = (m_prod * mi) // gcd(m_prod, mi)
return x % m_prod