Make prime factoring of perfect powers efficient

docs/pages/galois.rst

@@ -140,6 +140,8 @@ Modular Arithmetic
    lcm
    crt
    isqrt
+   iroot
+   ilog
    pow
    is_cyclic
    carmichael

galois/factor.py

@@ -8,13 +8,24 @@
 
 import numpy as np
 
-from .math_ import isqrt
+from .math_ import isqrt, iroot, ilog
 from .overrides import set_module
 from .prime import PRIMES, is_prime
 
 __all__ = ["prime_factors", "is_smooth"]
 
 
+def perfect_power(n):
+    max_prime_idx = bisect.bisect_right(PRIMES, ilog(n, 2))
+
+    for prime in PRIMES[0:max_prime_idx]:
+        x = iroot(n, prime)
+        if x**prime == n:
+            return x, prime
+
+    return None
+
+
 def trial_division_factor(n):
     max_factor = isqrt(n)
     max_prime_idx = bisect.bisect_right(PRIMES, max_factor)
@@ -76,8 +87,9 @@ def prime_factors(n):
     **Steps**:
 
     1. Test if :math:`n` is prime. If so, return `[n], [1]`.
-    2. Use trial division with a list of primes up to :math:`10^6`. If no residual factors, return the discovered prime factors.
-    3. Use Pollard's Rho algorithm to find a non-trivial factor of the residual. Continue until all are found.
+    2. Test if :math:`n` is a perfect power, such that :math:`n = x^k`. If so, prime factor :math:`x` and multiply its exponents by :math:`k`.
+    3. Use trial division with a list of primes up to :math:`10^6`. If no residual factors, return the discovered prime factors.
+    4. Use Pollard's Rho algorithm to find a non-trivial factor of the residual. Continue until all are found.
 
     Parameters
     ----------
@@ -122,9 +134,17 @@ def prime_factors(n):
         return [n], [1]
 
     # Step 2
-    p, e, n = trial_division_factor(n)
+    result = perfect_power(n)
+    if result is not None:
+        base, exponent = result
+        p, e = prime_factors(base)
+        e = [ei * exponent for ei in e]
+        return p, e
 
     # Step 3
+    p, e, n = trial_division_factor(n)
+
+    # Step 4
     while n > 1 and not is_prime(n):
         f = pollard_rho_factor(n)  # A non-trivial factor
         while f is None:

galois/math_.py

@@ -8,7 +8,7 @@
 
 from .overrides import set_module
 
-__all__ = ["pow", "isqrt", "lcm", "prod"]
+__all__ = ["pow", "isqrt", "iroot", "ilog", "lcm", "prod"]
 
 
 @set_module("galois")
@@ -72,11 +72,9 @@ def isqrt(n):
     --------
     .. ipython:: python
 
-        # Use a large Mersenne prime
-        p = galois.mersenne_primes(2000)[-1]; p
-        sqrt_p = galois.isqrt(p); sqrt_p
-        sqrt_p**2 <= p
-        (sqrt_p + 1)**2 <= p
+        galois.isqrt(27**2 - 1)
+        galois.isqrt(27**2)
+        galois.isqrt(27**2 + 1)
     """
     if sys.version_info.major == 3 and sys.version_info.minor >= 8:
         return math.isqrt(n)  # pylint: disable=no-member
@@ -98,6 +96,108 @@ def isqrt(n):
             return large_candidate
 
 
+@set_module("galois")
+def iroot(n, k):
+    """
+    Finds the integer :math:`k`-th root :math:`x` of :math:`n`, such that :math:`x^k \\le n`.
+
+    Parameters
+    ----------
+    n : int
+        A positive integer.
+    k : int
+        The root :math:`k`, must be at least 2.
+
+    Returns
+    -------
+    int
+        The integer :math:`k`-th root :math:`x` of :math:`n`, such that :math:`x^k \\le n`
+
+    Examples
+    --------
+    .. ipython :: python
+
+        galois.iroot(27**5 - 1, 5)
+        galois.iroot(27**5, 5)
+        galois.iroot(27**5 + 1, 5)
+    """
+    if not isinstance(n, (int, np.integer)):
+        raise TypeError(f"Argument `n` must be an integer, not {type(n)}.")
+    if not isinstance(k, (int, np.integer)):
+        raise TypeError(f"Argument `k` must be an integer, not {type(k)}.")
+    if not n > 0:
+        raise ValueError(f"Argument `n` must be positive, not {n}.")
+    if not k >= 2:
+        raise ValueError(f"Argument `k` must be at least 2, not {k}.")
+
+    # https://stackoverflow.com/a/39191163/11694321
+    u = n
+    x = n + 1
+    k1 = k - 1
+
+    while u < x:
+        x = u
+        u = (k1*u + n // u**k1) // k
+
+    return x
+
+
+@set_module("galois")
+def ilog(n, b):
+    """
+    Finds the integer :math:`\\textrm{log}_b(n) = k`, such that :math:`b^k \\le n`.
+
+    Parameters
+    ----------
+    n : int
+        A positive integer.
+    b : int
+        The logarithm base :math:`b`.
+
+    Returns
+    -------
+    int
+        The integer :math:`\\textrm{log}_b(n) = k`, such that :math:`b^k \\le n`.
+
+    Examples
+    --------
+    .. ipython :: python
+
+        galois.ilog(27**5 - 1, 27)
+        galois.ilog(27**5, 27)
+        galois.ilog(27**5 + 1, 27)
+    """
+    if not isinstance(n, (int, np.integer)):
+        raise TypeError(f"Argument `n` must be an integer, not {type(n)}.")
+    if not isinstance(b, (int, np.integer)):
+        raise TypeError(f"Argument `b` must be an integer, not {type(b)}.")
+    if not n > 0:
+        raise ValueError(f"Argument `n` must be positive, not {n}.")
+    if not b >= 2:
+        raise ValueError(f"Argument `b` must be at least 2, not {b}.")
+
+    # https://stackoverflow.com/a/39191163/11694321
+    low, b_low, high, b_high = 0, 1, 1, b
+
+    while b_high < n:
+        low, b_low, high, b_high = high, b_high, high*2, b_high**2
+
+    while high - low > 1:
+        mid = (low + high) // 2
+        b_mid = b_low * b**(mid - low)
+        if n < b_mid:
+            high, b_high = mid, b_mid
+        elif b_mid < n:
+            low, b_low = mid, b_mid
+        else:
+            return mid
+
+    if b_high == n:
+        return high
+
+    return low
+
+
 @set_module("galois")
 def lcm(*integers):
     """

tests/test_math.py

@@ -8,20 +8,55 @@
 
 def test_isqrt():
     p = galois.mersenne_primes(2000)[-1]
-
     sqrt_p = galois.isqrt(p)
     assert isinstance(sqrt_p, int)
     assert sqrt_p**2 <= p and not (sqrt_p + 1)**2 <= p
 
-    sqrt_p = galois.isqrt(p - 1)
+    p = galois.mersenne_primes(2000)[-1] - 1
+    sqrt_p = galois.isqrt(p)
     assert isinstance(sqrt_p, int)
     assert sqrt_p**2 <= p and not (sqrt_p + 1)**2 <= p
 
-    sqrt_p = galois.isqrt(p + 1)
+    p = galois.mersenne_primes(2000)[-1] + 1
+    sqrt_p = galois.isqrt(p)
     assert isinstance(sqrt_p, int)
     assert sqrt_p**2 <= p and not (sqrt_p + 1)**2 <= p
 
 
+def test_iroot():
+    p = galois.mersenne_primes(2000)[-1]
+    root = galois.iroot(p, 10)
+    assert isinstance(root, int)
+    assert root**10 <= p and not (root + 1)**10 <= p
+
+    p = galois.mersenne_primes(2000)[-1] - 1
+    root = galois.iroot(p, 10)
+    assert isinstance(root, int)
+    assert root**10 <= p and not (root + 1)**10 <= p
+
+    p = galois.mersenne_primes(2000)[-1] + 1
+    root = galois.iroot(p, 10)
+    assert isinstance(root, int)
+    assert root**10 <= p and not (root + 1)**10 <= p
+
+
+def test_ilog():
+    p = galois.mersenne_primes(2000)[-1]
+    exponent = galois.ilog(p, 17)
+    assert isinstance(exponent, int)
+    assert 17**exponent <= p and not 17**(exponent + 1) <= p
+
+    p = galois.mersenne_primes(2000)[-1] - 1
+    exponent = galois.ilog(p, 17)
+    assert isinstance(exponent, int)
+    assert 17**exponent <= p and not 17**(exponent + 1) <= p
+
+    p = galois.mersenne_primes(2000)[-1] + 1
+    exponent = galois.ilog(p, 17)
+    assert isinstance(exponent, int)
+    assert 17**exponent <= p and not 17**(exponent + 1) <= p
+
+
 def test_lcm():
     assert galois.lcm() == 1
     assert galois.lcm(1, 0, 4) == 0