Faster two squares

cp-algorithms · Oct 7, 2024 · 3687447 · 3687447
1 parent f00a116
commit 3687447
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/cp-algo/number_theory/two_squares.hpp b/cp-algo/number_theory/two_squares.hpp
@@ -39,7 +39,7 @@ namespace cp_algo::math {
             cnt[p]++;
         }
         // 1, -1, i, -i
-        std::vector<gaussint> res = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
+        std::vector<gaussint> res = {1};
         for(auto [p, c]: cnt) {
             std::vector<gaussint> nres;
             if(p % 4 == 3) {
@@ -67,8 +67,12 @@ namespace cp_algo::math {
         }
         std::vector<gaussint> nres;
         for(auto p: res) {
-            if(p.real() >= 0 && p.imag() >= 0) {
-                nres.push_back(p);
+            while(p.real() < 0 || p.imag() < 0) {
+                p *= gaussint(0, 1);
+            }
+            nres.push_back(p);
+            if(!p.real() || !p.imag()) {
+                nres.emplace_back(p.imag(), p.real());
             }
         }
         return nres;