from math import inf, sqrt
# Find (real) roots of a*x^2 + b*x + c.
# Returns the smallest non-negative solution first, followed by the
# other (smallest absolute value first, if there is no positive solution).
# If there are no real roots, or infinitely many (inf, inf) is returned.
# When dealing with a non-constant linear function its root and inf are
# returned.
def pqsolve(a, b, c, eps=1e-300):
if a == 0:
if b != 0: return -c/b, inf
else: return inf, inf
# Ensure that b > 0; by inverting the signs of
# of the coefficients the roots of the quadratic
# equation remain unchanged.
if b < 0:
a, b, c = -a, -b, -c
if b == 0:
if a*c <= 0: return sqrt(abs(c/a)), -sqrt(abs(c/a))
else: return inf, inf
else:
# If the discriminant is negative there are no real roots.
if 4 * a * c > b*b * (1 + eps):
return inf, inf
discriminant = max(0, b*b - 4 * a * c)
x1, x2 = (-(b + sqrt(discriminant)) / (2 * a),
-2 * c / (b + sqrt(discriminant)))
# Decide on the ordering, which is the
# "smaller" of the two roots.
if ((x2 >= 0 and x1 < 0)
or (x2 >= 0 and x2 < x1)
or (x2 <= 0 and x1 < x2)):
return x2, x1
return x1, x2
# Not numerically stable implementation of above function (pqsolve).
def pqsolve_naive(a, b, c):
if a == 0:
if b != 0: return -c/b, inf
else: return inf, inf
discriminant = b*b - 4 * a * c
x1, x2 = ((-b + sqrt(discriminant)) / (2 * a),
(-b - sqrt(discriminant)) / (2 * a))
if ((x2 >= 0 and x1 < 0)
or (x2 >= 0 and x2 < x1)
or (x2 <= 0 and x1 < x2)):
return x2, x1
return x1, x2
if __name__ == "__main__":
# None of these assertions should fail, if you implementend pqsolve correctly.
assert(pqsolve(1, 0, -1) == (1.0, -1.0))
assert(pqsolve(1, 0, 1) == (inf, inf))
assert(pqsolve(1e-300, 0, -1e-300) == (1.0, -1.0))
print("All good!")