Fix mod_sqrt for a = 1, p = 2

cheran-senthil · Nov 2, 2024 · f115dfc · f115dfc
1 parent d6d3efc
commit f115dfc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyrival/algebra/mod_sqrt.py b/pyrival/algebra/mod_sqrt.py
@@ -1,8 +1,8 @@
 def mod_sqrt(a, p):
     """returns x s.t. x**2 == a (mod p)"""
     a %= p
-    if a == 0:
-        return 0
+    if a < 2:
+        return a
     assert pow(a, (p - 1) // 2, p) == 1
     if p & 3 == 3:
         return pow(a, (p + 1) // 4, p)