Allow returning basis in solve

diophantine.py

@@ -40,7 +40,7 @@
 from builtins import next
 from builtins import range
 from copy import deepcopy
-from math import copysign, sqrt, log10, floor
+from math import sqrt, log10, floor
 try:
     from math import gcd  # Py >= 3.6
 except ImportError:
@@ -53,17 +53,12 @@ class NoSolutionException(Exception):
     pass
 
 
-# Sign of a variable, which isn't included in math for some reason
-def sign(x):
-    return copysign(1, x) if x else 0
-
-
 def nonzero(m):
     return [(i, j) for i, j in product(range(m.shape[0]), range(m.shape[1]))
             if m[i, j] != 0]
 
 
-def solve(A, b):
+def solve(A, b, return_basis=False):
     """
     Finds small solutions to systems of diophantine equations, A x = b, where A
     is a M x N matrix of coefficents, b is a M x 1 vector and x is the
@@ -110,6 +105,8 @@ def solve(A, b):
             raise NotImplementedError("Ax=B has unique solution in integers")
     else:
         solutions = []
+    if return_basis:
+        return solutions, basis
     return solutions
 
 
@@ -406,7 +403,7 @@ def addr(a, b, c, d):
 def comparer(a, b, c, d):
     """Assumes b>0 and d>0.  Returns -1, 0 or 1 according as a/b <,=,> c/d+ """
     assert b > 0 and d > 0
-    return sign(a * d - b * c)
+    return int(sign(a * d - b * c))
 
 
 def lcasvector(A, x):