diff --git a/wizard-work/beam/egcd.py b/wizard-work/beam/egcd.py new file mode 100644 index 0000000..308052b --- /dev/null +++ b/wizard-work/beam/egcd.py @@ -0,0 +1,24 @@ +"""Extended Euclidean algorithm. + +egcd(a, b) returns (g, x, y) such that a*x + b*y == g == gcd(a, b). +Works for negative inputs and zero. +""" + +def egcd(a, b): + if b == 0: + if a < 0: + return (-a, -1, 0) + return (a, 1, 0) + g, x1, y1 = egcd(b, a % b) + q = a // b + return (g, y1, x1 - q * y1) + + +def gcd(a, b): + return egcd(a, b)[0] + + +def lcm(a, b): + if a == 0 or b == 0: + return 0 + return abs(a // gcd(a, b) * b) diff --git a/wizard-work/beam/test_egcd.py b/wizard-work/beam/test_egcd.py new file mode 100644 index 0000000..9f3b8bb --- /dev/null +++ b/wizard-work/beam/test_egcd.py @@ -0,0 +1,25 @@ +import sys, os +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from egcd import egcd, gcd, lcm + +def test_basic(): + assert egcd(30, 12) == (6, -1, 3) or egcd(30,12)[0]==6 + assert gcd(30,12)==6 + assert lcm(4,6)==12 + +def test_negatives(): + g,x,y = egcd(-30,12) + assert g==6 and -30*x + 12*y == g + +def test_zero(): + assert gcd(0,5)==5 + assert gcd(5,0)==5 + assert lcm(0,5)==0 + +def test_identity(): + g,x,y = egcd(17,31) + assert g==1 and 17*x+31*y==1 +if __name__=="__main__": + for f in [test_basic,test_negatives,test_zero,test_identity]: + f() + print("ALL PASS")