Ticket #6509: trac_6509_four_squares.patch

sage/rings/arith.py

@@ -3773 +3773 @@
     else:
         raise RuntimeError, "unknown algorithm '%s'"%algorithm
 
+def _brute_force_four_squares(n):
+    """
+    Brute force search for decomposition into a sum of four squares,
+    for cases that the main algorithm fails to handle.
+
+    INPUT: a positive integer
+    OUTPUT: a list of four numbers whose squares sum to n
+
+    EXAMPLES::
+        sage: from sage.rings.arith import _brute_force_four_squares
+        sage: _brute_force_four_squares(567)
+        [1, 1, 6, 23]
+    """
+    from math import sqrt
+    for i in range(0,int(sqrt(n))+1):
+        for j in range(i,int(sqrt(n-i**2))+1):
+            for k in range(j, int(sqrt(n-i**2-j**2))+1):
+                rem = n-i**2-j**2-k**2
+                if rem >= 0:
+                    l = int(sqrt(rem))
+                    if rem-l**2==0:
+                        return [i,j,k,l]
+    
+def four_squares(n):
+    """
+    Computes the decomposition into the sum of four squares,
+    using an algorithm described by Peter Schorn at:
+    http://www.schorn.ch/howto.html.
+
+    INPUT: an integer
+
+    OUTPUT: a list of four numbers whose squares sum to n
+
+    EXAMPLES::
+        sage: four_squares(3)
+        [0, 1, 1, 1]
+        sage: four_squares(130)
+        [0, 0, 3, 11]
+        sage: four_squares(1101011011004)
+        [2, 1049178, 2370, 15196]
+        sage: sum([i-sum([q^2 for q in four_squares(i)]) for i in range(2,10000)]) # long time
+        0
+    """
+    from sage.rings.integer_mod import mod
+    from math import sqrt
+    from sage.rings.arith import _brute_force_four_squares
+    try: 
+        ts = two_squares(n)
+        return [0,0,ts[0],ts[1]]
+    except ValueError:
+        pass
+    m = n
+    v = 0
+    while mod(m,4)==0:
+        v = v +1
+        m = m/4
+    if mod(m,8) == 7:
+        d = 1
+        m = m - 1
+    else:
+        d = 0
+    if mod(m,8)==3:
+        x = int(sqrt(m))
+        if mod(x,2) == 0:
+            x = x - 1
+        p = (m-x**2)/2
+        while not is_prime(p):
+            x = x - 2
+            p = (m-x**2)/2
+            if x < 0:
+            # fall back to brute force
+                m = m + d
+                return [2**v*q for q in _brute_force_four_squares(m)]
+        y,z = two_squares(p)
+        return [2**v*q for q in [d,x,y+z,abs(y-z)]]
+    x = int(sqrt(m))
+    p = m - x**2
+    if p == 1:
+        return[2**v*q for q in [d,0,x,1]]
+    while not is_prime(p):
+        x = x - 1
+        p = m - x**2
+        if x < 0:
+            # fall back to brute force
+            m = m + d
+            return [2**v*q for q in _brute_force_four_squares(m)]
+    y,z = two_squares(p)
+    return [2**v*q for q in [d,x,y,z]]
+
+
 def subfactorial(n):
     r"""
     Subfactorial or rencontres numbers, or derangements: number of