update many_facts

adamant-pwn · adamant-pwn · commit d131d012a51a · 2025-05-26T20:24:56.000+02:00
diff --git a/verify/simd/many_facts.test.cpp b/verify/simd/many_facts.test.cpp
@@ -2,9 +2,9 @@
 #define PROBLEM "https://judge.yosupo.jp/problem/many_factorials"
 #pragma GCC optimize("Ofast,unroll-loops")
 #include <bits/stdc++.h>
-#define CP_ALGO_CHECKPOINT
-#include "cp-algo/util/checkpoint.hpp"
+//#define CP_ALGO_CHECKPOINT
 #include "blazingio/blazingio.min.hpp"
+#include "cp-algo/util/checkpoint.hpp"
 #include "cp-algo/util/simd.hpp"
 #include "cp-algo/math/common.hpp"
 
@@ -14,84 +14,97 @@ using namespace cp_algo;
 constexpr int mod = 998244353;
 constexpr int imod = -math::inv2(mod);
 
-void facts_inplace(vector<int> &args) {
-    constexpr int block = 1 << 16;
-    static basic_string<size_t> args_per_block[mod / block];
-    uint64_t limit = 0;
-    for(auto [i, x]: args | views::enumerate) {
-        if(x < mod / 2) {
-            limit = max(limit, uint64_t(x));
-            args_per_block[x / block].push_back(i);
-        } else {
-            limit = max(limit, uint64_t(mod - x - 1));
-            args_per_block[(mod - x - 1) / block].push_back(i);
+vector<int> facts(vector<int> const& args) {
+    constexpr int accum = 4;
+    constexpr int simd_size = 8;
+    constexpr int block = 1 << 18;
+    constexpr int subblock = block / simd_size;
+    static basic_string<array<int, 2>> odd_args_per_block[mod / subblock];
+    static basic_string<array<int, 2>> reg_args_per_block[mod / subblock];
+    constexpr int limit_reg = mod / 64;
+    int limit_odd = 0;
+
+    vector<int> res(size(args), 1);
+    auto prod_mod = [&](uint64_t a, uint64_t b) {
+        return (a * b) % mod;
+    };
+    for(auto [i, xy]: views::zip(args, res) | views::enumerate) {
+        auto [x, y] = xy;
+        auto t = x;
+        if(t >= mod / 2) {
+            t = mod - t - 1;
+            y = t % 2 ? 1 : mod - 1;
+        }
+        int pw = 0;
+        while(t > limit_reg) {
+            limit_odd = max(limit_odd, (t - 1) / 2);
+            odd_args_per_block[(t - 1) / 2 / subblock].push_back({int(i), (t - 1) / 2});
+            t /= 2;
+            pw += t;
         }
+        reg_args_per_block[t / subblock].push_back({int(i), t});
+        y = int(y * math::bpow(2, pw, 1ULL, prod_mod) % mod);
     }
     cp_algo::checkpoint("init");
     uint32_t b2x32 = (1ULL << 32) % mod;
-    uint64_t fact = 1;
-    const int accum = 4;
-    const int simd_size = 8;
-    for(uint64_t b = 0; b <= limit; b += accum * block) {
-        u32x8 cur[accum];
-        static array<u32x8, block / simd_size> prods[accum];
-        for(int z = 0; z < accum; z++) {
-            for(int j = 0; j < simd_size; j++) {
-                cur[z][j] = uint32_t(b + z * block + j * (block / simd_size));
-                prods[z][0][j] = cur[z][j] + !(b || z || j);
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
-                cur[z][j] = uint32_t(uint64_t(cur[z][j]) * b2x32 % mod);
-#pragma GCC diagnostic pop
-            }
-        }
-        for(int i = 1; i < block / simd_size; i++) {
+    auto process = [&](int limit, auto &args_per_block, auto step, auto &&proj) {
+        uint64_t fact = 1;
+        for(int b = 0; b <= limit; b += accum * block) {
+            u32x8 cur[accum];
+            static array<u32x8, subblock> prods[accum];
             for(int z = 0; z < accum; z++) {
-                cur[z] += b2x32;
-                cur[z] = cur[z] >= mod ? cur[z] - mod : cur[z];
-                prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod, imod);
-            }
-        }
-        cp_algo::checkpoint("inner loop");
-        for(int z = 0; z < accum; z++) {
-            uint64_t bl = b + z * block;
-            for(auto i: args_per_block[bl / block]) {
-                size_t x = args[i];
-                if(x >= mod / 2) {
-                    x = mod - x - 1;
-                }
-                x -= bl;
-                auto pre_blocks = x / (block / simd_size);
-                auto in_block = x % (block / simd_size);
-                auto ans = fact * prods[z][in_block][pre_blocks] % mod;
-                for(size_t j = 0; j < pre_blocks; j++) {
-                    ans = ans * prods[z].back()[j] % mod;
+                for(int j = 0; j < simd_size; j++) {
+                    cur[z][j] = uint32_t(b + z * block + j * subblock);
+                    cur[z][j] = proj(cur[z][j]);
+                    prods[z][0][j] = cur[z][j] + !cur[z][j];
+    #pragma GCC diagnostic push
+    #pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+                    cur[z][j] = uint32_t(uint64_t(cur[z][j]) * b2x32 % mod);
+    #pragma GCC diagnostic pop
                 }
-                if(args[i] >= mod / 2) {
-                    ans = math::bpow(ans, mod - 2, 1ULL, [](auto a, auto b){return a * b % mod;});
-                    args[i] = int(x % 2 ? ans : mod - ans);
-                } else {
-                    args[i] = int(ans);
+            }
+            for(int i = 1; i < block / simd_size; i++) {
+                for(int z = 0; z < accum; z++) {
+                    cur[z] += step;
+                    cur[z] = cur[z] >= mod ? cur[z] - mod : cur[z];
+                    prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod, imod);
                 }
             }
-            args_per_block[bl / block].clear();
-            for(int j = 0; j < simd_size; j++) {
-                fact = fact * prods[z].back()[j] % mod;
+            cp_algo::checkpoint("inner loop");
+            for(int z = 0; z < accum; z++) {
+                for(int j = 0; j < simd_size; j++) {
+                    int bl = b + z * block + j * subblock;
+                    for(auto [i, x]: args_per_block[bl / subblock]) {
+                        auto ans = fact * prods[z][x - bl][j] % mod;
+                        res[i] = int(res[i] * ans % mod);
+                    }
+                    fact = fact * prods[z].back()[j] % mod;
+                }
             }
+            cp_algo::checkpoint("mul ans");
+        }
+    };
+    uint32_t b2x33 = (1ULL << 33) % mod;
+    process(limit_reg, reg_args_per_block, b2x32, identity{});
+    process(limit_odd, odd_args_per_block, b2x33, [](uint32_t x) {return 2 * x + 1;});
+    for(auto [i, x]: res | views::enumerate) {
+        if (args[i] >= mod / 2) {
+            x = int(math::bpow(x, mod - 2, 1ULL, prod_mod));
         }
-        cp_algo::checkpoint("write ans");
     }
+    cp_algo::checkpoint("inv ans");
+    return res;
 }
 
 void solve() {
     int n;
     cin >> n;
     vector<int> args(n);
     for(auto &x : args) {cin >> x;}
-    cp_algo::checkpoint("input read");
-    facts_inplace(args);
-    for(auto it: args) {cout << it << "\n";}
-    cp_algo::checkpoint("output written");
+    cp_algo::checkpoint("read");
+    auto res = facts(args);
+    for(auto it: res) {cout << it << "\n";}
+    cp_algo::checkpoint("write");
     cp_algo::checkpoint<1>();
 }
 
