#! /usr/bin/python3 # -*- coding: UTF-8 -*- # Formatted for tab-stops 4 # # By Peter Schorn, peter.schorn@acm.org, December 2004 - July 2020 # # Decompose a non-negative integer into the sum of at most four squares. # Last change 24-Jul-2020, adapted for Python 3 # Last change 25-Nov-2005, simplified decomposeProbablePrime # Last change 13-Nov-2005, optimized search loop # Last change 30-Oct-2005, improved isqrt function # Last change 8-May-2005, improved isqrt function import sys, operator, math, random from functools import reduce # Return the greatest common divisor using Euclid's algorithm def gcd(a, b): while b: a, b = b, a % b return a # 'l' is a list of tuples where the first elements form the argument to 'func' # and the last element is the expected result. Returns True iff all tests # succeed. def testfunc(func, l): return False not in [func(*x[:-1]) == x[-1] for x in l] # Test the gcd function # Returns True if no error has been detected def testgcd(): return testfunc(gcd, [(4, 6, 2), (6, 4, 2), (1, 1, 1), (2, 1, 1), (1, 2, 1), (13000, 17000, 1000), (3, 4, 1)]) # Compute the jacobi symbol (a / b) # Precondition: b odd # Postcondition: Result is jacobi symbol (a / b) or 0 iff gcd(a, b) != 1 def jacobi(a, b): s = 1 while a > 1: d4 = a >> 2 m4 = a & 3 if m4: # a % 4 != 0 if m4 == 2: # a % 4 = 2 (i.e. a % 2 = 0) if b & 7 in [3, 5]: s = -s a >>= 1 else: if m4 == b & 3 == 3: s = -s a, b = b % a, a else: # a % 4 = 0 a = d4 # a = a / 4 return a and s # Test the jacobi function by comparing with a pre-computed table # Returns True if no error has been detected def testjacobi(): return testfunc(jacobi, [ (1, 1, 1), (1, 3, 1), (1, 5, 1), (1, 7, 1), (1, 9, 1), (1, 11, 1), (1, 13, 1), (1, 15, 1), (1, 17, 1), (1, 19, 1), (1, 21, 1), (1, 23, 1), (1, 25, 1), (1, 27, 1), (1, 29, 1), (1, 31, 1), (2, 1, 1), (2, 3, -1), (2, 5, -1), (2, 7, 1), (2, 9, 1), (2, 11, -1), (2, 13, -1), (2, 15, 1), (2, 17, 1), (2, 19, -1), (2, 21, -1), (2, 23, 1), (2, 25, 1), (2, 27, -1), (2, 29, -1), (2, 31, 1), (3, 1, 1), (3, 3, 0), (3, 5, -1), (3, 7, -1), (3, 9, 0), (3, 11, 1), (3, 13, 1), (3, 15, 0), (3, 17, -1), (3, 19, -1), (3, 21, 0), (3, 23, 1), (3, 25, 1), (3, 27, 0), (3, 29, -1), (3, 31, -1), (4, 1, 1), (4, 3, 1), (4, 5, 1), (4, 7, 1), (4, 9, 1), (4, 11, 1), (4, 13, 1), (4, 15, 1), (4, 17, 1), (4, 19, 1), (4, 21, 1), (4, 23, 1), (4, 25, 1), (4, 27, 1), (4, 29, 1), (4, 31, 1), (5, 1, 1), (5, 3, -1), (5, 5, 0), (5, 7, -1), (5, 9, 1), (5, 11, 1), (5, 13, -1), (5, 15, 0), (5, 17, -1), (5, 19, 1), (5, 21, 1), (5, 23, -1), (5, 25, 0), (5, 27, -1), (5, 29, 1), (5, 31, 1), (6, 1, 1), (6, 3, 0), (6, 5, 1), (6, 7, -1), (6, 9, 0), (6, 11, -1), (6, 13, -1), (6, 15, 0), (6, 17, -1), (6, 19, 1), (6, 21, 0), (6, 23, 1), (6, 25, 1), (6, 27, 0), (6, 29, 1), (6, 31, -1), (7, 1, 1), (7, 3, 1), (7, 5, -1), (7, 7, 0), (7, 9, 1), (7, 11, -1), (7, 13, -1), (7, 15, -1), (7, 17, -1), (7, 19, 1), (7, 21, 0), (7, 23, -1), (7, 25, 1), (7, 27, 1), (7, 29, 1), (7, 31, 1), (8, 1, 1), (8, 3, -1), (8, 5, -1), (8, 7, 1), (8, 9, 1), (8, 11, -1), (8, 13, -1), (8, 15, 1), (8, 17, 1), (8, 19, -1), (8, 21, -1), (8, 23, 1), (8, 25, 1), (8, 27, -1), (8, 29, -1), (8, 31, 1), (9, 1, 1), (9, 3, 0), (9, 5, 1), (9, 7, 1), (9, 9, 0), (9, 11, 1), (9, 13, 1), (9, 15, 0), (9, 17, 1), (9, 19, 1), (9, 21, 0), (9, 23, 1), (9, 25, 1), (9, 27, 0), (9, 29, 1), (9, 31, 1), (10, 1, 1), (10, 3, 1), (10, 5, 0), (10, 7, -1), (10, 9, 1), (10, 11, -1), (10, 13, 1), (10, 15, 0), (10, 17, -1), (10, 19, -1), (10, 21, -1), (10, 23, -1), (10, 25, 0), (10, 27, 1), (10, 29, -1), (10, 31, 1), (11, 1, 1), (11, 3, -1), (11, 5, 1), (11, 7, 1), (11, 9, 1), (11, 11, 0), (11, 13, -1), (11, 15, -1), (11, 17, -1), (11, 19, 1), (11, 21, -1), (11, 23, -1), (11, 25, 1), (11, 27, -1), (11, 29, -1), (11, 31, -1), (12, 1, 1), (12, 3, 0), (12, 5, -1), (12, 7, -1), (12, 9, 0), (12, 11, 1), (12, 13, 1), (12, 15, 0), (12, 17, -1), (12, 19, -1), (12, 21, 0), (12, 23, 1), (12, 25, 1), (12, 27, 0), (12, 29, -1), (12, 31, -1), (13, 1, 1), (13, 3, 1), (13, 5, -1), (13, 7, -1), (13, 9, 1), (13, 11, -1), (13, 13, 0), (13, 15, -1), (13, 17, 1), (13, 19, -1), (13, 21, -1), (13, 23, 1), (13, 25, 1), (13, 27, 1), (13, 29, 1), (13, 31, -1), (14, 1, 1), (14, 3, -1), (14, 5, 1), (14, 7, 0), (14, 9, 1), (14, 11, 1), (14, 13, 1), (14, 15, -1), (14, 17, -1), (14, 19, -1), (14, 21, 0), (14, 23, -1), (14, 25, 1), (14, 27, -1), (14, 29, -1), (14, 31, 1), (15, 1, 1), (15, 3, 0), (15, 5, 0), (15, 7, 1), (15, 9, 0), (15, 11, 1), (15, 13, -1), (15, 15, 0), (15, 17, 1), (15, 19, -1), (15, 21, 0), (15, 23, -1), (15, 25, 0), (15, 27, 0), (15, 29, -1), (15, 31, -1)]) # Return True iff n is a prime. Algorithm uses trial division and is therefore # only usable for 'small' n. This function is used to test the probable prime # functions which work for 'large' n as well. def isPrime(n): if n <= 1: return False if (n <= 3) or (n == 5): return True if (n & 1 == 0) or (n % 3 == 0): return False t = 5 d = 2 while (t * t <= n) and (n % t): t += d d = 6 - d # skip trial divisors t with t % 3 = 0 return n % t != 0 # Low quality primality test # Precondition: n is odd # For n > 2: if result is False then n is composite # if result is True then n maybe prime # 341 (= 11 * 31) is smallest composite number where function returns True def isProbablePrime2(n): return (n > 1) and (pow(2, n - 1, n) == 1) # Test the probable prime detector 'func' on all odd numbers between # n and upperN. Return False iff there is a case where a prime number # is erroneously flagged as composite. def testisProbablePrimeFunc(func, n, upperN): if n & 1 == 0: n += 1 # Make sure that n is odd while n <= upperN: if func(n) < isPrime(n): # Error if composite was prime after all return False n += 2 # Next odd return True # The list of all odd primes less than 100 oddPrimes97 = list(filter(isProbablePrime2, list(range(100)))) # The product of all odd primes less than 100 oddPrimeProduct97 = reduce(operator.mul, oddPrimes97) # Good quality primality test # Precondition: n is odd # For n > 2: if result is False then n is composite # if result if True then n maybe prime # 42799 (= 127 * 337) is smallest composite number where function returns True # Idea is to check the identity jacobi(2, p) = 2 ^ ((p - 1) / 2) % p # which holds for primes (according to the Euler criterion) # Note that isProbablePrime(1093 * 1093) = True def isProbablePrime(n): return (n > 1) and ((n < 561) or gcd(oddPrimeProduct97, n) == 1) and \ (pow(2, n >> 1, n) == ((n & 7 in [1, 7]) or (n - 1))) # Use the fact that True == 1 # Return True iff n is 2 or a strong pseudo prime to base a # Checking for strong pseudo primality is the core of the Miller-Rabin # probabilistic prime check: if a number is a strong pseudo prime for # sufficiently many bases, it is highly likely to be a prime. Note that this # idea does not work with just checking a^(n-1) = 1 mod n as a Carmichael # number c (such as 561) will satisfy this criterion for all a with gcd(a,c) = 1 def isStrongPseudoPrime(a, n): if n == 2: return True n1 = n - 1 # i.e. -1 % n r = n1 >> 1 k = 1 q = r >> 1 m = r & 1 while m == 0: k += 1 r = q q = r >> 1 m = r & 1 # Postcondition: (n = 2^k * r + 1) and (r % 2 = 1) t = pow(a, r, n) pt = n1 # pt will hold previous value of t while (t != 1) and (t != n1) and k: pt = t # save previous value t = (t * t) % n # square modulo n k -= 1 # We have a strong pseudo prime iff # Case 1) We reached 1 and the previous value was -1 (mod n) or # Case 2) We reached -1 and we have at least one squaring to go return (t == 1) and (pt == n1) or (t == n1) and (k > 0) # Note: We cannot replace (k > 0) with k as a boolean value is # expected. primeProduct3 = 5*7*13*23*29*31*37*53*61*97*127*149*151*157*137*229 # Very good quality primality test (but slower than isProbablePrime) # For n > 0: if result is False then n is composite or 1 # if result if True then n maybe prime # 220729 (= 103 * 2143) is smallest composite number where function returns True def isProbablePrime3(n): return (n > 1) and ((n < 2047) or gcd(primeProduct3, n) == 1) and \ isStrongPseudoPrime(2, n) # Test the probable prime detector 'func' on all odd numbers between # n and upperN. Return False iff there is a discrepancy between # 'func' and isPrime def testisProbablePrimeFuncStrict(func, n, upperN): if n & 1 == 0: n += 1 # Make sure that n is odd while n <= upperN: if func(n) != isPrime(n): # Error if functions disagree print('Prime error for %i. Got %s which is not correct.' \ % (n, func(n))) return False n += 2 # Next odd return True # Probable prime check to be used. Note that correctness of algorithms in this # module only depends on the fact that composite numbers are detected correctly. useIsProbablePrime = isProbablePrime # productive use #useIsProbablePrime = isProbablePrime # productive use #useIsProbablePrime = isProbablePrime2 # testing #useIsProbablePrime = lambda x: x > 1 # testing #useIsProbablePrime = isProbablePrime3 # testing # Performance evaluation # Case 'small': time decompose.py 0 100000 f # Case 'medium': time decompose.py 'sys.maxint-500L' 100000 f # Case 'large': time decompose.py '10**200' 50 f # # (Mac OS X 10.3.8, 1.5GHz MHz PowerPC G4, 256 KB L2 Cache, 2MB L3 Cache, # 1.5GB Memory, Python 2.3 (#1, Sep 13 2003, 00:49:11) # [GCC 3.3 20030304 (Apple Computer, Inc. build 1495)] on darwin ) # small large total # isProbablePrime 0:11.47 0:05.45 16.92 # isProbablePrime2 0:10.08 0:27.89 37.97 # lambda x: x > 1 0:10.90 0:27.73 38.63 # isProbablePrime3 0:13.44 0:13.66 27.10 # Cache for the odd primes less than 500 oddPrimesCache = list(filter(isProbablePrime, list(range(500)))) # Return the factorial of n def fac(n): return reduce(operator.mul, list(range(1, n + 1)), 1) # Return the product of the first n primes # pp(0) = 1, pp(1) = 2, pp(2) = 6, ... def pp(n): if n > 0: return reduce(operator.mul, [p for p in oddPrimes(n - 1)], 1) << 1 return 1 def testpp(): return testfunc(pp, [(-10, 1), (-1, 1), (0, 1), (1, 2), (2, 6), (3, 30), (9, 223092870)]) # Return the smallest probable prime p>=n with p % (m/gcd(m,r))=(r/gcd(m,r)) def dp(n, m, r): g = gcd(m, r) m //= g p = m * (n // m) + r // g if p < n: p += m while not isProbablePrime(p): p += m return p def testdp(): return testfunc(dp, [(100, 4, 1, 101), (101, 4, 1, 101), (100, 6, 4, 101), (101, 6, 4, 101)]) # Return next probable prime >= n # Precondition: n is odd def nextProbablePrime(n): while not useIsProbablePrime(n): n += 2 return n # Return next probable prime >= n def np(n): if n <= 2: return 2 return nextProbablePrime((n & 1) and n or (n + 1)) # For given even k and odd prime p, return True if n = k * p + 1 is prime. # If True is returned, n is guaranteed to be prime but note that the result # might be False even though n is prime. def isnpk(k, p): n = k * p + 1 if gcd(n, oddPrimeProduct97) > 1: return n in oddPrimes97 # in case n is a prime less than 100 t = pow(2, k, n) if t == 1: return False return pow(t, p, n) == 1 def testisnpk(): return testfunc(isnpk, [(4, 13, True), (2, 5003, True), (4, 5003, False), (2, (1 << 31) - 1, False), (46, (1 << 31) - 1, True), (32, 3, True), (20, 5, True), (22, 5, False)]) # For given k and odd prime p return K >= k such that n = K * p + 1 is prime. # Note that n is then guaranteed to be prime but K is not necessarily the # smallest value >= k to make this happen. Note also that Dirichlet's theorem # guarantees the existence of such a prime (but this algorithm might fail # for all such primes ...) def npk(k, p): if k & 1: k += 1 # k is now even while not isnpk(k, p): k += 2 return k # Generator for the odd primes def oddPrimes(limit = -1): index = 0 while (limit < 0) or (index < limit): if index >= len(oddPrimesCache): oddPrimesCache.append(nextProbablePrime(oddPrimesCache[-1] + 2)) yield oddPrimesCache[index] index += 1 # Probabilistic check whether n is a square. If the result is (False, q) then n # cannot be a square and (q / n) = -1. If the result is (True, 0), then n may # or may not be a square. # The idea is to compute the jacobi symbol (q / n) for 'certainty' many odd # primes q. Since (q / n1 * n2) = (q / n1) * (q / n2), the jacobi symbol (q / n) # is always 1 if n is a square. This also means that if we can find a q such # that (q / n) = -1 then n cannot be a square. # Heuristically each additional q filters out above half of the non-squares. # The number n = 182919 is the smallest non-square for which # isSquareProbabilistic(n, 15) returns True, 0 # The number n = 8042479 is the smallest non-square for which # isSquareProbabilistic(n, 30) returns True, 0 def isSquareProbabilistic(n, certainty): generator = oddPrimes() while certainty: q = next(generator) if jacobi(n % q, q) == -1: # q is a witness that n is not a square return False, q certainty -= 1 return True, 0 # There are exactly 336 squares mod 10080. This is the smallest rate of squares # for all n < 15840. magicN = 10080 squaresModMagicN = {}.fromkeys([(i * i) % magicN for i in range(magicN >> 1)]) # If result is True then n may or may not be a square. If the result is False # then n cannot be a square. The 'false positive' rate is 3.3% = 336 / 10080 = # 1 / 30. The smallest non-square for which isProbableSquare(n) = True is 385. def isProbableSquare(n): return (n % magicN) in squaresModMagicN def testisProbableSquare(low, high): while low <= high: if not isProbableSquare(low) and (isqrt(low)[1] == 0): return False low += 1 return True # Compute an imaginary unit modulo p # Precondition: p % 4 = 1 # Postcondition: p is prime implies: Result is (x, False) and x^2 % p = -1 # The result can also be (x, True) and then x^2 = p # Note that this algorithm might succeed even when the argument is not prime. # Example (1) # n = 3277 = 29 * 113 % 8 = 5, 2^((n-1)/4) = 2^819 = 128 mod n # 128^2 % n = -1 # Explanation: # order(2) mod 29 = order(2) mod 113 = order(2) mod n = 28 divides n - 1 # This implies: 2^(n-1) mod 29 = 2^(n-1) mod 113 = 1 # We also have 2^((n-1)/2) mod 29 = 2^((n-1)/2) mod 113 = -1 # overall 2^((n-1)/2) % n = -1 # Example (2) # p = 3281 = 17 * 193 % 8 = 1, jacobi(3, p) = -1 # 3^((p-1)/4) = 3^820 = 81 mod p, 81^2 = -1 mod p # Example (3) # p = 49141 = 157 * 313 % 8 = 5, 2^((p-1)/4) = 2^12285 = 32527 mod p # 32527^2 = -1 mod p, isProbablePrime(p) = True # Example (4) # p = 665281 = 577 * 1153 % 8 = 1, jacobi(13, p) = -1 # 13 ^ ((p-1)/4) = 13 ^ 166320 = 531393 mod p # 531393^2 = -1 mod p, isProbablePrime(p) = True # Note also that if p is a square, jacobi(q, p) = 1 for all q def iunit(p): if p & 7 == 5: q = 2 else: primes = oddPrimes() q = next(primes) # q is an odd (probable) prime # (q % 2 = 1) and (p % 4 = 1) implies jacobi(q, p) = jacobi(p % q, q) while jacobi(p % q, q) == 1: # jacobi(q, p) q = next(primes) if (q == 229) and isProbableSquare(p):# loop is running quite long # reached for p = dp(1, 4*pp(k), 1) for k > 47 or p = 1093 ** 2 s, r = isqrt(p) # check, if p is a square if r == 0: # if yes, return square root return s, True return pow(q, p >> 2, p), False # Test the iunit function for arguments between p and upperP making sure # that the argument is a (probable) prime congruent 1 modulo 4 # Returns True if no error has been detected def testiunit(p, upperP): p = ((p >> 2) << 2) + 1 # Make sure that p % 4 = 1 upperP = min(upperP, 42799) # and that upperP does not exceed range where while p <= upperP: # isProbablePrime works correctly. if isProbablePrime(p): if (iunit(p)[0] ** 2 + 1) % p: # This is the check return False p += 4 # maintain invariant p % 4 = 1 return True ln2 = math.log(2) # Returns the smallest integer greater or equal to the logarithm of n to base 2, # i.e. the result r satisfies 2^r <= n < 2^(r+1). def floorLog2(n): r = int(0.5 + math.log(n) / ln2) return ((1 << r) <= n) and r or (r - 1) # The square root algorithm uses its iterative variant for numbers less than # 2^1024. For larger numbers the recursive case is invoked. isqrtIterativeBetter = 1 << 1024 # Input: n >= 0, log2n = floor(log(2, n)), 2^log2n <= n < 2^(log2n + 1) # Output: (s, r) where (s-1)^2 <= n < (s+1)^2 and s^2 + r = n # For very small n, the built-in square root function is invoked. For larger n # (typically n < 10^300) an iterative algorithm is used while for the largest # n, a divide & conquer approach is employed almost identical to the one from # Paul Zimmermann (Karatsuba Square Root, https://hal.inria.fr/inria-00072854) # The difference is that a slightly weaker invariant ((s-1)^2 <= n < (s+1)^2 # can be shown to suffice - note that the r = n - s^2 can be negative in the # case that the square root is over estimated. def isqrtInternal(n, log2n): if n <= sys.maxsize: s = int(math.sqrt(n)) return s, n - s * s if n < isqrtIterativeBetter: d = 7 * (log2n // 14 - 1) q = 7 # assert (0 < (n >> (d + d)) < (1 << 28)) and (q <= d) s = int(math.sqrt(n >> (d + d))) while d: # assert (s - 1) ** 2 <= (n >> (d + d)) < (s + 1) ** 2 if q > d: q = d s <<= q d -= q q += q s = (s + (n >> (d + d)) // s) >> 1 return s, n - s * s log2b = log2n >> 2 mask = (1 << log2b) - 1 s, r = isqrtInternal(n >> (log2b + log2b), log2n - (log2b + log2b)) q, u = divmod((r << log2b) + ((n >> log2b) & mask), s + s) return (s << log2b) + q, (u << log2b) + (n & mask) - q * q # Input: n >= 0 # Output: (s, r) where s^2 <= n < (s+1)^2 and s^2 + r = n def isqrt(n): assert n >= 0 if n <= sys.maxsize: s = int(math.sqrt(n)) return s, n - s * s s, r = isqrtInternal(n, floorLog2(n)) return (r < 0) and (s - 1, r + s + s - 1) or (s, r) # Test the isqrt function for arguments between n and upperN # Returns True if no error has been detected def testisqrt(n, upperN): while n <= upperN: s, r = isqrt(n) if not ((s * s <= n < (s + 1) ** 2) and (s * s + r == n)): return False # return False iff postcondition of isqrt not met n += 1 return True # Try to decompose n into a sum of two squares # Precondition: n is a probable prime, n % 4 = 1 # (a, b, True) with a^2 + b^2 = n iff decomposition was successful # (0, 0, False) iff decomposition was not successful def decomposeProbablePrime(n): b, r = iunit(n) if r: return 0, b, True if (b * b + 1) % n: # Check whether we got an imaginary unit return 0, 0, False # Indicate failure if not a = n while b * b > n: a, b = b, a % b return b, a % b, True # Perform n tests starting at the smallest prime p >= start with p % 4 = 1 def testdecomposeProbablePrime(start, n): p = ((start >> 2) << 2) + 1 if p < start: p += 4 while n: while not isProbablePrime(p): p += 4 a, b, c = decomposeProbablePrime(p) if c and (a * a + b * b != p): return False n -= 1 p += 4 return True # Dictionary of exceptional cases the search in decompose cannot handle decomposeExceptions = { 9634: [56, 57, 57], 2986: [21, 32, 39], 1906: [13, 21, 36], 1414: [ 6, 17, 33], 730: [ 0, 1, 27], 706: [15, 15, 16], 526: [ 6, 7, 21], 370: [ 8, 9, 15], 226: [ 8, 9, 9], 214: [ 3, 6, 13], 130: [ 0, 3, 11], 85: [ 0, 6, 7], 58: [ 0, 3, 7], 34: [ 3, 3, 4], 10: [ 0, 1, 3], 3: [ 1, 1, 1], 2: [ 0, 1, 1] } # Decompose an integer n>=0 into a sum of up to four squares. Result is a # tuple where the first entry is a list of four integers and the second # entry is True iff the search loop was successful. Note that currently # there is no known input where the search loop would not be successful # but there is no known proof. def decompose(n): # Multiply all elements of l with v and return a tuple with the sorted list # and the value True indicating success def sortScaled(l): l.sort() return [v * x for x in l], True assert n >= 0 if n <= 1: # Done if n = 0 or n = 1 return [0, 0, 0, n], True v = 1 # Remove powers of 4 while n & 3 == 0: # n % 4 = 0 n >>= 2 # n = n / 4 v += v # Postcondition: (nOld = v^2*n) and (n%4 != 0) and (v = 2^k for some k >= 0) sq, p = isqrt(n) if p == 0: # n is a square return [0, 0, 0, v * sq], True if (n & 3 == 1) and useIsProbablePrime(n): # n is prime, n % 4 = 1 s, r, done = decomposeProbablePrime(n) if done: # otherwise n was not a prime return sortScaled([0, 0, s, r]) if n & 7 == 7: # Need 4 squares, subtract largest square delta^2 # such that (n > delta^2) and (delta^2 % 8 != 0) delta, n = sq, p if sq & 3 == 0: delta -= 1 n += delta + delta + 1 # sq % 4 = 0 -> delta % 8 in [3, 7] -> # delta^2 % 8 = 1 -> n % 8 = 6 # sq % 4 != 0 -> delta % 8 in [1, 2, 3, 5, 6, 7] -> # delta^2 % 8 in [1, 4] -> n % 8 in [3, 6] # and this implies n % 4 != 1, i.e. n cannot be a sum of two squares sq, p = isqrt(n) # sq^2 != n since n % 4 = 3 which is never true for a square else: delta = 0 # Postcondition: (sq = isqrt(n)) and (n % 8 != 7) and (n % 4 != 0) and # (nOld = v^2 * (n + delta^2)) # This implies that n is a sum of three squares - now check whether n # is one of the special cases the rest of the algorithm could not handle. if n in decomposeExceptions: # Retrieve pre-computed result and scale with v return sortScaled([delta] + decomposeExceptions[n]) # Now perform search distinguishing two cases noting that n % 4 != 0 # Case 1: n % 4 = 3, n = x^2 + 2*p, p % 4 = 1, # p prime, p = y^2 + z^2 implies n = x^2 + (y + z)^2 + (y - z)^2 if n & 3 == 3: if sq & 1 == 0: sq -= 1 p += sq + sq + 1 p >>= 1 while True: if useIsProbablePrime(p): r0, r1, done = decomposeProbablePrime(p) if done: return sortScaled([delta, sq, r0 + r1, abs(r0 - r1)]) sq -= 2 # Next sq if sq < 0: # Should not happen, no known case return 4 * [0], False p += sq + sq + 2 # Next p # Case 2: n % 4 = 1 or n % 4 = 2, n = x^2 + p, # p % 4 = 1, p prime, p = y^2 + z^2 implies n = x^2 + y^2 + z^2 if (n - sq) & 1 == 0: sq -= 1 p += sq + sq + 1 while True: if useIsProbablePrime(p): r0, r1, done = decomposeProbablePrime(p) if done: return sortScaled([delta, sq, r0, r1]) sq -= 2 # Next sq if sq < 0: # Should not happen, no known case return 4 * [0], False p += 4 * (sq + 1) # Next p # Decompose all integers from list li into at most four squares and print # the results iff verbose = True. The result is a list of all cases where the # decomposition was not successful. Note that this list should be empty. def checkDecompose(li, verbose): f = [] for i in li: r, done = decompose(i) if verbose: # Use the fact that r is sorted if r[0]: # Need four squares s = '%i = %i^2 + %i^2 + %i^2 + %i^2' % (i, r[0], r[1], r[2], r[3]) elif r[1]: # Need three squares s = '%i = %i^2 + %i^2 + %i^2' % (i, r[1], r[2], r[3]) elif r[2]: # Need two squares s = '%i = %i^2 + %i^2' % (i, r[2], r[3]) else: # Number is a square s = '%i = %i^2' % (i, r[3]) # print also number of squares needed print('%s # %i' % (s, 4 - r[:-1].count(0))) if not done or (sum([x*x for x in r]) != i): f.append((i, r, done)) return f # Perform a small self test of the major function. # Return True iff all tests successful. def selfTest(): return (testgcd() and testisProbablePrimeFunc(isProbablePrime2, 1, 1000) and testisProbablePrimeFunc(isProbablePrime, 1, 1000) and testjacobi() and testisqrt(0, 125) and testisqrt(10**20 + random.randint(0, 10**19), 20) and testiunit(5, 2377) and (iunit(1093 ** 2)[0] == 1093) and testisProbablePrimeFuncStrict(isProbablePrime3, 3, 1000) and testisProbablePrimeFuncStrict(isProbablePrime, 3, 1000) and testisnpk() and testpp() and testdecomposeProbablePrime(5, 100) and testdecomposeProbablePrime(10 ** 5 + 1, 3) and testisProbableSquare(100, 400) and (len(checkDecompose([dp(1, 4 * pp(48), 1)], False)) == 0) and testdp() ) def main(): # Print usage information for the script and exit. def usage(): print('''\ Usage: %s start(Exp) [iterations(Exp)] [v(char)] [f(char)] Decomposes the integers in start (if start evaluates to a list) or in [start, start + iterations - 1] (if start evaluates to an integer) into the sum of at most 4 squares. Expression can contain the following functions: gcd(a,b): greatest common divisor of a and b jacobi(a,b): the Jacobi symbol of a and b fac(n): factorial of n pp(n): product of the first n primes [pp(0)=1, pp(1)=2, pp(2)=6, ...] dp(n,m,r): smallest probable prime p >= n with p mod m/gcd(m,r)=r/gcd(m,r) np(n): smallest probable prime p >= n isqrt(n)[0]: integer square root of n is y with y^2 <= n < (y + 1)^2 v: verbose f: fast (omit self test) Examples: decompose.py 10007 decompose.py '577*1153' decompose.py 'np(19081961*10**3)' 2 v decompose.py '[2*10**i+19081961 for i in [34, 35, 36]]' v time decompose.py 0 100000 ''' % sys.argv[0]) sys.exit(2) # Called because some form of command line syntax error detected # Return the value of 'intString' if it contains a valid expression. # Otherwise print 'errorMessage', the usage information and exit. def getInt(intString, errorMessage): try: return eval(intString) except: print(errorMessage % intString) usage() # Return True iff string 'opt' is in the lower case version # of the script's arguments. def hasOption(opt): return opt in [x.lower() for x in sys.argv] if not (2 <= len(sys.argv) <= 5): usage() if not (hasOption('f') or selfTest()): print('Self test failed! Exiting...') sys.exit(1) li = getInt(sys.argv[1], 'Illegal expression \'%s\' for start.') verbose = hasOption('v') if not isinstance(li, list): # Create list of numbers to check. # If number of iterations is 0, use 1 instead. li = list(range(li, li + ((len(sys.argv) > 2) and not sys.argv[2].lower() in ['f', 'v'] and getInt(sys.argv[2], 'Illegal expression \'%s\' for number of iterations.') or 1))) r = checkDecompose(li, verbose or (len(li) == 1)) if len(r): print('Errors: ', r) if __name__ == '__main__': main()