Tiny simplification

cheran-senthil · Jul 9, 2023 · 7f54a87 · 7f54a87
1 parent dc477c8
commit 7f54a87
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/pyrival/algebra/fft.py b/pyrival/algebra/fft.py
@@ -13,9 +13,9 @@ def fft(P):
 
     k = n
     while k > 1:
-        for i in range(0, n, k):
-            r = rt[i//k]
-            for j1 in range(i, i + k//2):
+        for i in range(n//k):
+            r = rt[i]
+            for j1 in range(i*k, i*k + k//2):
                 j2 = j1 + k//2
                 z = r * P[j2]
                 P[j2] = P[j1] - z

diff --git a/pyrival/algebra/ntt.py b/pyrival/algebra/ntt.py
@@ -21,9 +21,9 @@ def ntt(P):
 
     k = n
     while k > 1:
-        for i in range(0, n, k):
-            r = rt[i//k]
-            for j1 in range(i, i + k//2):
+        for i in range(n//k):
+            r = rt[i]
+            for j1 in range(i*k, i*k + k//2):
                 j2 = j1 + k//2
                 z = r * P[j2]
                 P[j2] = (P[j1] - z) % MOD