Optimize PrimitiveRoot

kzrnm · Jan 17, 2023 · b4174db · b4174db
1 parent 1c3eacd
commit b4174db
Showing 1 changed file with 11 additions and 12 deletions.
diff --git a/Source/AtCoderLibrary/Math/Internal/InternalMath.cs b/Source/AtCoderLibrary/Math/Internal/InternalMath.cs
@@ -1,6 +1,9 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
+#if NETCOREAPP3_1_OR_GREATER
+using System.Numerics;
+#endif
 
 namespace AtCoder.Internal
 {
@@ -46,26 +49,22 @@ public static int PrimitiveRoot<TMod>() where TMod : struct, IStaticMod
         }
         static int PrimitiveRootCalculate<TMod>() where TMod : struct, IStaticMod
         {
-            int m = (int)default(TMod).Mod;
-            Span<int> divs = stackalloc int[20];
+            var m = default(TMod).Mod;
+            Span<uint> divs = stackalloc uint[20];
             divs[0] = 2;
             int cnt = 1;
-            int x = (m - 1) / 2;
-
-            while (x % 2 == 0)
-            {
-                x >>= 1;
-            }
+            var x = m - 1;
+            x >>= BitOperations.TrailingZeroCount(x);
 
-            for (int i = 3; (long)i * i <= x; i += 2)
+            for (uint i = 3; (long)i * i <= x; i += 2)
             {
                 if (x % i == 0)
                 {
                     divs[cnt++] = i;
-                    while (x % i == 0)
+                    do
                     {
                         x /= i;
-                    }
+                    } while (x % i == 0);
                 }
             }
 
@@ -81,7 +80,7 @@ static int PrimitiveRootCalculate<TMod>() where TMod : struct, IStaticMod
                     if (new StaticModInt<TMod>(g).Pow((m - 1) / d).Value == 1)
                         goto NEXT;
                 return g;
-            NEXT: { }
+            NEXT:;
             }
         }