Refactor PrimeFactorization: replace trial-division isPrime with Miller-Rabin, add tests (#1312)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: williamfiset

src/main/java/com/williamfiset/algorithms/math/PrimeFactorization.java

@@ -1,66 +1,136 @@
+/**
+ * Prime factorization using Pollard's rho algorithm with Miller-Rabin primality testing.
+ *
+ * <p>Miller-Rabin with the deterministic witness set {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}
+ * is guaranteed correct for all n < 3.317 × 10^24, which covers the full range of Java's long.
+ *
+ * <p>Time: O(n^(1/4)) expected per factor found
+ *
+ * @author William Fiset, william.alexandre.fiset@gmail.com
+ */
 package com.williamfiset.algorithms.math;
 
 import java.util.ArrayList;
-import java.util.PriorityQueue;
+import java.util.List;
 
 public class PrimeFactorization {
 
-  public static ArrayList<Long> primeFactorization(long n) {
-    ArrayList<Long> factors = new ArrayList<>();
-    if (n <= 0) throw new IllegalArgumentException();
-    else if (n == 1) return factors;
-    PriorityQueue<Long> divisorQueue = new PriorityQueue<>();
-    divisorQueue.add(n);
-    while (!divisorQueue.isEmpty()) {
-      long divisor = divisorQueue.remove();
-      if (isPrime(divisor)) {
-        factors.add(divisor);
-        continue;
-      }
-      long next_divisor = pollardRho(divisor);
-      if (next_divisor == divisor) {
-        divisorQueue.add(divisor);
-      } else {
-        divisorQueue.add(next_divisor);
-        divisorQueue.add(divisor / next_divisor);
-      }
-    }
+  /**
+   * Returns the prime factorization of n. The returned factors are not necessarily sorted.
+   *
+   * @throws IllegalArgumentException if n <= 0.
+   */
+  public static List<Long> primeFactorization(long n) {
+    if (n <= 0)
+      throw new IllegalArgumentException();
+
+    List<Long> factors = new ArrayList<>();
+    factor(n, factors);
     return factors;
   }
 
+  private static void factor(long n, List<Long> factors) {
+    if (n == 1)
+      return;
+    if (isPrime(n)) {
+      factors.add(n);
+      return;
+    }
+    long d = pollardRho(n);
+    factor(d, factors);
+    factor(n / d, factors);
+  }
+
   private static long pollardRho(long n) {
-    if (n % 2 == 0) return 2;
+    if (n % 2 == 0)
+      return 2;
     long x = 2 + (long) (999999 * Math.random());
     long c = 2 + (long) (999999 * Math.random());
     long y = x;
     long d = 1;
     while (d == 1) {
-      x = (x * x + c) % n;
-      y = (y * y + c) % n;
-      y = (y * y + c) % n;
+      x = mulMod(x, x, n) + c;
+      if (x >= n)
+        x -= n;
+      y = mulMod(y, y, n) + c;
+      if (y >= n)
+        y -= n;
+      y = mulMod(y, y, n) + c;
+      if (y >= n)
+        y -= n;
       d = gcd(Math.abs(x - y), n);
-      if (d == n) break;
+      if (d == n)
+        break;
     }
     return d;
   }
 
-  private static long gcd(long a, long b) {
-    return b == 0 ? a : gcd(b, a % b);
-  }
+  /**
+   * Deterministic Miller-Rabin primality test, correct for all long values. Uses 12 witnesses that
+   * guarantee correctness for n < 3.317 × 10^24.
+   */
+  private static boolean isPrime(long n) {
+    if (n < 2)
+      return false;
+    if (n < 4)
+      return true;
+    if (n % 2 == 0 || n % 3 == 0)
+      return false;
+
+    // Write n-1 as 2^r * d
+    long d = n - 1;
+    int r = Long.numberOfTrailingZeros(d);
+    d >>= r;
 
-  private static boolean isPrime(final long n) {
-    if (n < 2) return false;
-    if (n == 2 || n == 3) return true;
-    if (n % 2 == 0 || n % 3 == 0) return false;
-    long limit = (long) Math.sqrt(n);
-    for (long i = 5; i <= limit; i += 6) if (n % i == 0 || n % (i + 2) == 0) return false;
+    for (long a : new long[]{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}) {
+      if (a >= n)
+        continue;
+      long x = powMod(a, d, n);
+      if (x == 1 || x == n - 1)
+        continue;
+      boolean composite = true;
+      for (int i = 0; i < r - 1; i++) {
+        x = mulMod(x, x, n);
+        if (x == n - 1) {
+          composite = false;
+          break;
+        }
+      }
+      if (composite)
+        return false;
+    }
     return true;
   }
 
+  /** Modular exponentiation: (base^exp) % mod, using mulMod to avoid overflow. */
+  private static long powMod(long base, long exp, long mod) {
+    long result = 1;
+    base %= mod;
+    while (exp > 0) {
+      if ((exp & 1) == 1)
+        result = mulMod(result, base, mod);
+      exp >>= 1;
+      base = mulMod(base, base, mod);
+    }
+    return result;
+  }
+
+  /** Overflow-safe modular multiplication: (a * b) % mod. */
+  private static long mulMod(long a, long b, long mod) {
+    return java.math.BigInteger.valueOf(a)
+        .multiply(java.math.BigInteger.valueOf(b))
+        .mod(java.math.BigInteger.valueOf(mod))
+        .longValue();
+  }
+
+  private static long gcd(long a, long b) {
+    return b == 0 ? a : gcd(b, a % b);
+  }
+
   public static void main(String[] args) {
-    System.out.println(primeFactorization(7)); // [7]
-    System.out.println(primeFactorization(100)); // [2,2,5,5]
-    System.out.println(primeFactorization(666)); // [2,3,3,37]
-    System.out.println(primeFactorization(872342345)); // [5, 7, 7, 67, 19, 2797]
+    System.out.println(primeFactorization(7));         // [7]
+    System.out.println(primeFactorization(100));       // [2, 2, 5, 5]
+    System.out.println(primeFactorization(666));       // [2, 3, 3, 37]
+    System.out.println(primeFactorization(872342345)); // [5, 7, 7, 19, 67, 2797]
   }
 }

src/test/java/com/williamfiset/algorithms/math/BUILD

@@ -26,6 +26,17 @@ java_test(
     deps = TEST_DEPS,
 )
 
+# bazel test //src/test/java/com/williamfiset/algorithms/math:PrimeFactorizationTest
+java_test(
+    name = "PrimeFactorizationTest",
+    srcs = ["PrimeFactorizationTest.java"],
+    main_class = "org.junit.platform.console.ConsoleLauncher",
+    use_testrunner = False,
+    args = ["--select-class=com.williamfiset.algorithms.math.PrimeFactorizationTest"],
+    runtime_deps = JUNIT5_RUNTIME_DEPS,
+    deps = TEST_DEPS,
+)
+
 # bazel test //src/test/java/com/williamfiset/algorithms/math:EulerTotientFunctionTest
 java_test(
     name = "EulerTotientFunctionTest",

src/test/java/com/williamfiset/algorithms/math/PrimeFactorizationTest.java

@@ -0,0 +1,90 @@
+package com.williamfiset.algorithms.math;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.*;
+import org.junit.jupiter.api.*;
+
+public class PrimeFactorizationTest {
+
+  private static void assertFactors(long n, Long... expected) {
+    List<Long> factors = PrimeFactorization.primeFactorization(n);
+    Collections.sort(factors);
+    assertThat(factors).containsExactlyElementsIn(Arrays.asList(expected)).inOrder();
+  }
+
+  @Test
+  public void testOne() {
+    assertFactors(1);
+  }
+
+  @Test
+  public void testSmallPrimes() {
+    assertFactors(2, 2L);
+    assertFactors(3, 3L);
+    assertFactors(5, 5L);
+    assertFactors(7, 7L);
+    assertFactors(11, 11L);
+    assertFactors(13, 13L);
+  }
+
+  @Test
+  public void testPowersOfTwo() {
+    assertFactors(4, 2L, 2L);
+    assertFactors(8, 2L, 2L, 2L);
+    assertFactors(1024, 2L, 2L, 2L, 2L, 2L, 2L, 2L, 2L, 2L, 2L);
+  }
+
+  @Test
+  public void testPrimePowers() {
+    assertFactors(27, 3L, 3L, 3L);
+    assertFactors(125, 5L, 5L, 5L);
+    assertFactors(49, 7L, 7L);
+  }
+
+  @Test
+  public void testComposites() {
+    assertFactors(100, 2L, 2L, 5L, 5L);
+    assertFactors(666, 2L, 3L, 3L, 37L);
+    assertFactors(872342345, 5L, 7L, 7L, 19L, 67L, 2797L);
+  }
+
+  @Test
+  public void testProductEqualsOriginal() {
+    long[] values = {12, 360, 872342345, 1000000007, 239821585064027L};
+    for (long n : values) {
+      List<Long> factors = PrimeFactorization.primeFactorization(n);
+      long product = 1;
+      for (long f : factors)
+        product *= f;
+      assertThat(product).isEqualTo(n);
+    }
+  }
+
+  @Test
+  public void testLargePrime() {
+    assertFactors(8763857775536878331L, 8763857775536878331L);
+  }
+
+  @Test
+  public void testLargeSemiprime() {
+    // 15485867 and 15486481 are both prime
+    assertFactors(239821585064027L, 15485867L, 15486481L);
+  }
+
+  @Test
+  public void testLargeComposite() {
+    assertFactors(1000000007L * 999999937L, 999999937L, 1000000007L);
+  }
+
+  @Test
+  public void testZeroThrows() {
+    assertThrows(IllegalArgumentException.class, () -> PrimeFactorization.primeFactorization(0));
+  }
+
+  @Test
+  public void testNegativeThrows() {
+    assertThrows(IllegalArgumentException.class, () -> PrimeFactorization.primeFactorization(-5));
+  }
+}