218 lines
8.1 KiB
Python
218 lines
8.1 KiB
Python
"""Exact integer arithmetic: primes, divisibility, modular arithmetic, combinatorics, bases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Union
|
|
|
|
import sympy as sp
|
|
from sympy.ntheory.modular import crt
|
|
|
|
from mcp_calculator.ops.common import MathError, check_list, clip, to_str
|
|
from mcp_calculator.ops.parsing import parse
|
|
from mcp_calculator.ops.spec import NumberTheoryOperation
|
|
|
|
IntLike = Union[int, str]
|
|
|
|
# operation -> (min values, max values, value names for error messages)
|
|
_ARITY: dict[str, tuple[int, int, str]] = {
|
|
"gcd": (2, 10_000, "[a, b, ...]"),
|
|
"lcm": (2, 10_000, "[a, b, ...]"),
|
|
"factorize": (1, 1, "[n]"),
|
|
"is_prime": (1, 1, "[n]"),
|
|
"next_prime": (1, 1, "[n]"),
|
|
"previous_prime": (1, 1, "[n]"),
|
|
"nth_prime": (1, 1, "[k]"),
|
|
"primes_in_range": (2, 2, "[low, high]"),
|
|
"divisors": (1, 1, "[n]"),
|
|
"totient": (1, 1, "[n]"),
|
|
"mod": (2, 2, "[a, m]"),
|
|
"mod_pow": (3, 3, "[base, exponent, modulus]"),
|
|
"mod_inverse": (2, 2, "[a, m]"),
|
|
"crt": (2, 10_000, "[r1, r2, ...] plus 'moduli'"),
|
|
"factorial": (1, 1, "[n]"),
|
|
"binomial": (2, 2, "[n, k]"),
|
|
"permutations": (2, 2, "[n, k]"),
|
|
"fibonacci": (1, 1, "[n]"),
|
|
"continued_fraction": (1, 1, "[x] (rational or expression)"),
|
|
"base_convert": (1, 1, "[value] plus 'from_base' / 'to_base'"),
|
|
}
|
|
|
|
_MAX_FACTORIAL = 20_000
|
|
_MAX_PRIMES = 10_000
|
|
_DIGITS = "0123456789abcdefghijklmnopqrstuvwxyz"
|
|
|
|
|
|
def number_theory(
|
|
operation: NumberTheoryOperation,
|
|
values: list[IntLike],
|
|
moduli: list[IntLike] | None = None,
|
|
from_base: int = 10,
|
|
to_base: int = 10,
|
|
) -> dict[str, Any]:
|
|
if operation not in _ARITY:
|
|
raise MathError(f"unknown operation {operation!r}")
|
|
lo, hi, names = _ARITY[operation]
|
|
check_list(values, "values", min_items=lo, max_items=hi)
|
|
if len(values) > hi or len(values) < lo:
|
|
raise MathError(f"'{operation}' expects values = {names}")
|
|
|
|
if operation == "base_convert":
|
|
return _base_convert(str(values[0]), from_base, to_base)
|
|
if operation == "continued_fraction":
|
|
x = parse(str(values[0]))
|
|
if not x.is_real:
|
|
raise MathError("continued fractions need a real number")
|
|
max_terms = 50 if x.is_Rational else 20
|
|
terms = _continued_fraction(x, max_terms)
|
|
out: dict[str, Any] = {"result": terms}
|
|
if len(terms) == max_terms:
|
|
out["note"] = f"first {max_terms} terms"
|
|
return out
|
|
|
|
n = [_int(v) for v in values]
|
|
match operation:
|
|
case "gcd":
|
|
result: Any = sp.gcd_list(n)
|
|
case "lcm":
|
|
result = sp.lcm_list(n)
|
|
case "factorize":
|
|
_check_magnitude(n[0], 10**60, "factorize")
|
|
factors = sp.factorint(n[0])
|
|
pretty = " * ".join(f"{p}^{e}" if e > 1 else str(p) for p, e in factors.items()) or str(n[0])
|
|
return {"result": {str(p): e for p, e in factors.items()}, "product_form": pretty}
|
|
case "is_prime":
|
|
result = bool(sp.isprime(n[0]))
|
|
case "next_prime":
|
|
result = sp.nextprime(n[0])
|
|
case "previous_prime":
|
|
if n[0] <= 2:
|
|
raise MathError("there is no prime below 2")
|
|
result = sp.prevprime(n[0])
|
|
case "nth_prime":
|
|
if not 1 <= n[0] <= 10**7:
|
|
raise MathError("k must be between 1 and 10^7")
|
|
result = sp.prime(n[0])
|
|
case "primes_in_range":
|
|
low, high = n
|
|
primes = list(sp.primerange(low, high + 1)) if high - low <= 10**7 else None
|
|
if primes is None:
|
|
raise MathError("range too large (max width 10^7)")
|
|
out = {"result": primes[:_MAX_PRIMES], "count": len(primes)}
|
|
if len(primes) > _MAX_PRIMES:
|
|
out["note"] = f"only the first {_MAX_PRIMES} are listed"
|
|
return out
|
|
case "divisors":
|
|
_check_magnitude(n[0], 10**40, "divisors")
|
|
divs = sp.divisors(n[0])
|
|
return {"result": divs[:_MAX_PRIMES], "count": len(divs), "sum": sum(divs)}
|
|
case "totient":
|
|
_check_magnitude(n[0], 10**60, "totient")
|
|
result = sp.totient(n[0])
|
|
case "mod":
|
|
if n[1] == 0:
|
|
raise MathError("modulus must not be 0")
|
|
result = n[0] % n[1]
|
|
case "mod_pow":
|
|
if n[2] == 0:
|
|
raise MathError("modulus must not be 0")
|
|
try:
|
|
result = pow(n[0], n[1], n[2])
|
|
except ValueError:
|
|
raise MathError("a negative exponent needs base and modulus to be coprime") from None
|
|
case "mod_inverse":
|
|
try:
|
|
result = sp.mod_inverse(n[0], n[1])
|
|
except ValueError:
|
|
raise MathError(f"{n[0]} has no inverse modulo {n[1]} (not coprime)") from None
|
|
case "crt":
|
|
if not moduli or len(moduli) != len(n):
|
|
raise MathError("'crt' needs 'moduli' with one modulus per residue in 'values'")
|
|
m = [_int(v) for v in moduli]
|
|
solution = crt(m, n)
|
|
if solution is None:
|
|
return {"result": None, "note": "no solution: the congruences are inconsistent"}
|
|
return {"result": int(solution[0]), "modulus": int(solution[1]), "meaning": f"x ≡ {solution[0]} (mod {solution[1]})"}
|
|
case "factorial":
|
|
if not 0 <= n[0] <= _MAX_FACTORIAL:
|
|
raise MathError(f"n must be between 0 and {_MAX_FACTORIAL}")
|
|
result = sp.factorial(n[0])
|
|
case "binomial":
|
|
result = sp.binomial(n[0], n[1])
|
|
case "permutations":
|
|
if n[1] > n[0] or n[1] < 0:
|
|
result = 0
|
|
else:
|
|
result = sp.ff(n[0], n[1])
|
|
case "fibonacci":
|
|
if not 0 <= n[0] <= 100_000:
|
|
raise MathError("n must be between 0 and 100000")
|
|
result = sp.fibonacci(n[0])
|
|
case _:
|
|
raise MathError(f"unknown operation {operation!r}")
|
|
|
|
return _int_out(result)
|
|
|
|
|
|
def _int_out(result: Any) -> dict[str, Any]:
|
|
if isinstance(result, bool):
|
|
return {"result": result}
|
|
value = int(result)
|
|
text = str(value)
|
|
out: dict[str, Any] = {"result": clip(text) if len(text) > 15 else value}
|
|
if len(text) > 15:
|
|
out["digits"] = len(text.lstrip("-"))
|
|
out["note"] = "integer given as string to keep every digit"
|
|
return out
|
|
|
|
|
|
def _int(value: IntLike) -> int:
|
|
if isinstance(value, bool):
|
|
raise MathError("booleans are not integers here")
|
|
if isinstance(value, int):
|
|
return value
|
|
if isinstance(value, float):
|
|
if not value.is_integer():
|
|
raise MathError(f"{value} is not an integer")
|
|
return int(value)
|
|
parsed = parse(str(value))
|
|
if not parsed.is_Integer:
|
|
raise MathError(f"{value!r} is not an integer")
|
|
return int(parsed)
|
|
|
|
|
|
def _check_magnitude(n: int, bound: int, operation: str) -> None:
|
|
if abs(n) > bound:
|
|
raise MathError(f"'{operation}' is limited to |n| <= 10^{len(str(bound)) - 1}")
|
|
|
|
|
|
def _continued_fraction(x: sp.Expr, max_terms: int) -> list[int]:
|
|
terms = []
|
|
for term in sp.continued_fraction_iterator(x):
|
|
terms.append(int(term))
|
|
if len(terms) >= max_terms:
|
|
break
|
|
return terms
|
|
|
|
|
|
def _base_convert(value: str, from_base: int, to_base: int) -> dict[str, Any]:
|
|
if not (2 <= from_base <= 36 and 2 <= to_base <= 36):
|
|
raise MathError("bases must be between 2 and 36")
|
|
text = value.strip().lower().replace("_", "")
|
|
for prefix, base in (("0x", 16), ("0b", 2), ("0o", 8)):
|
|
if text.lstrip("-").startswith(prefix) and from_base == base:
|
|
text = text.replace(prefix, "", 1)
|
|
if len(text) > 1000:
|
|
raise MathError("value too long")
|
|
try:
|
|
n = int(text, from_base)
|
|
except ValueError:
|
|
raise MathError(f"{value!r} is not a valid base-{from_base} number") from None
|
|
if n == 0:
|
|
digits = "0"
|
|
else:
|
|
digits, m = "", abs(n)
|
|
while m:
|
|
m, r = divmod(m, to_base)
|
|
digits = _DIGITS[r] + digits
|
|
digits = ("-" if n < 0 else "") + digits
|
|
return {"result": digits, "decimal": to_str(n), "from_base": from_base, "to_base": to_base}
|