diff --git a/blockchain/diophantine_equation.py b/blockchain/diophantine_equation.py index ae6a145d2922..5fc3777c431e 100644 --- a/blockchain/diophantine_equation.py +++ b/blockchain/diophantine_equation.py @@ -100,6 +100,64 @@ def extended_gcd(a: int, b: int) -> tuple[int, int, int]: return (d, x, y) +def all_diophantine_solutions( + a: int, + b: int, + c: int, + n: int = 2, +) -> list[tuple[int, int]]: + """ + Return up to `n` integer solutions (x, y) to the linear Diophantine equation + a*x + b*y = c using the extended Euclidean algorithm. + + Raises + ------ + ValueError + If no integer solutions exist. + + Time complexity + --------------- + O(log(max(|a|, |b|))) to compute a base solution using extended_gcd; + plus O(n) to enumerate `n` solutions. + + Space complexity + ---------------- + O(1) beyond the returned list. + + Examples + -------- + >>> all_diophantine_solutions(10, 6, 14, n=2) + [(-7, 14), (-4, 9)] + >>> all_diophantine_solutions(10, 6, 14, n=4) + [(-7, 14), (-4, 9), (-1, 4), (2, -1)] + >>> all_diophantine_solutions(3, 6, 10, n=1) + Traceback (most recent call last): + ... + ValueError: No integer solutions exist for a=3, b=6, c=10 + """ + if a == 0 and b == 0: + if c == 0: + # Infinite solutions; return one canonical solution. + return [(0, 0)][: min(1, n)] + raise ValueError("No integer solutions exist for a=0, b=0, c!=0") + + g, xg, yg = extended_gcd(abs(a), abs(b)) + if c % g != 0: + msg = f"No integer solutions exist for a={a}, b={b}, c={c}" + raise ValueError(msg) + + # Scale a particular solution to ax + by = c + x0, y0 = xg * (c // g), yg * (c // g) + if a < 0: + x0 = -x0 + if b < 0: + y0 = -y0 + + # General solution: x = x0 + t*(b/g), y = y0 - t*(a/g) + dx, dy = b // g, a // g + return [(x0 + t * dx, y0 - t * dy) for t in range(n)] + + if __name__ == "__main__": from doctest import testmod