79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
"""Modular exponentiation and modular inverse.
|
|
|
|
Public API:
|
|
modpow(base, exp, mod) -> int
|
|
Compute base**exp mod mod efficiently (O(log exp)).
|
|
Supports negative exp when base is invertible mod mod (returns inverse).
|
|
mod must be a positive integer. Raises ValueError otherwise.
|
|
|
|
modinv(a, mod) -> int
|
|
Modular inverse of a modulo mod (a*result ≡ 1 mod mod).
|
|
Raises ValueError if gcd(a, mod) != 1.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
|
|
def _validate_mod(mod):
|
|
if not isinstance(mod, int) or isinstance(mod, bool):
|
|
raise TypeError("mod must be an int")
|
|
if mod <= 0:
|
|
raise ValueError("mod must be positive")
|
|
|
|
|
|
def modpow(base, exp, mod):
|
|
_validate_mod(mod)
|
|
if not isinstance(exp, int) or isinstance(exp, bool):
|
|
raise TypeError("exp must be an int")
|
|
if not isinstance(base, int) or isinstance(base, bool):
|
|
raise TypeError("base must be an int")
|
|
base %= mod
|
|
if mod == 1:
|
|
return 0
|
|
if exp < 0:
|
|
# base must be invertible
|
|
g, x, _ = _egcd(base, mod)
|
|
if g != 1:
|
|
raise ValueError("base is not invertible mod mod")
|
|
base = x % mod
|
|
exp = -exp
|
|
result = 1 % mod
|
|
b = base % mod
|
|
e = exp
|
|
while e > 0:
|
|
if e & 1:
|
|
result = (result * b) % mod
|
|
b = (b * b) % mod
|
|
e >>= 1
|
|
return result
|
|
|
|
|
|
def _egcd(a, b):
|
|
# iterative extended gcd
|
|
old_r, r = a, b
|
|
old_s, s = 1, 0
|
|
old_t, t = 0, 1
|
|
while r != 0:
|
|
q = old_r // r
|
|
old_r, r = r, old_r - q * r
|
|
old_s, s = s, old_s - q * s
|
|
old_t, t = t, old_t - q * t
|
|
return old_r, old_s, old_t
|
|
|
|
|
|
def modinv(a, mod):
|
|
_validate_mod(mod)
|
|
if not isinstance(a, int) or isinstance(a, bool):
|
|
raise TypeError("a must be an int")
|
|
if mod == 1:
|
|
return 0
|
|
g, x, _ = _egcd(a % mod, mod)
|
|
if g != 1:
|
|
raise ValueError("a is not invertible mod mod")
|
|
return x % mod
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(modpow(2, 10, 1000)) # 1024 % 1000 = 24
|
|
print(modinv(3, 11)) # 4
|
|
print(modpow(3, -1, 11)) # 4
|