support multivariate convolution (truncated)

Author: adamant-pwn

cp-algo/math/fft.hpp

@@ -29,6 +29,7 @@ namespace cp_algo::math::fft {
             }
         }
 
+        dft(size_t n): A(n), B(n) {init();}
         dft(auto const& a, size_t n): A(n), B(n) {
             init();
             base b2x32 = bpow(base(2), 32);
@@ -69,8 +70,8 @@ namespace cp_algo::math::fft {
                 B.fft();
             }
         }
-
-        void dot(auto &&C, auto const& D) {
+        template<bool overwrite = true>
+        void dot(auto const& C, auto const& D, auto &Aout, auto &Bout, auto &Cout) const {
             cvector::exec_on_evals<1>(A.size() / flen, [&](size_t k, point rt) {
                 k *= flen;
                 auto [Ax, Ay] = A.at(k);
@@ -93,13 +94,23 @@ namespace cp_algo::math::fft {
                     real(Dv)[0] = dx * real(rt) - dy * imag(rt);
                     imag(Dv)[0] = dx * imag(rt) + dy * real(rt);
                 }
-                A.at(k) = AC;
-                C.at(k) = AD + BC;
-                B.at(k) = BD;
+                if(overwrite) {
+                    Aout.at(k) = AC;
+                    Cout.at(k) = AD + BC;
+                    Bout.at(k) = BD;
+                } else {
+                    Aout.at(k) += AC;
+                    Cout.at(k) += AD + BC;
+                    Bout.at(k) += BD;
+                }
             });
             checkpoint("dot");
         }
 
+        void dot(auto &&C, auto const& D) {
+            dot(C, D, A, B, C);
+        }
+
         void recover_mod(auto &&C, auto &res, size_t k) {
             size_t check = (k + flen - 1) / flen * flen;
             assert(res.size() >= check);

cp-algo/math/multivar.hpp

@@ -0,0 +1,86 @@
+#ifndef CP_ALGO_MATH_MULTIVAR_HPP
+#define CP_ALGO_MATH_MULTIVAR_HPP
+#include "../util/big_alloc.hpp"
+#include "../number_theory/modint.hpp"
+#include "../math/fft.hpp"
+namespace cp_algo::math::fft {
+    template<modint_type base>
+    struct multivar {
+        std::vector<base, cp_algo::big_alloc<base>> data;
+        std::vector<size_t, cp_algo::big_alloc<size_t>> ranks;
+        std::vector<size_t> dim;
+        size_t N;
+        size_t rank(size_t i) {
+            size_t cur = 1, res = 0, K = size(dim);
+            for(auto ni: dim) {
+                cur *= ni;
+                res += i / cur;
+            }
+            return res % K;
+        }
+        multivar(std::vector<size_t> const& dim): dim(dim), N(
+            std::ranges::fold_left(dim, 1, std::multiplies{})
+        ) {
+            data.resize(N);
+            ranks.resize(N);
+            for(auto [i, x]: ranks | std::views::enumerate) {
+                x = rank(i);
+            }
+        }
+        void read() {
+            for(auto &it: data) {
+                std::cin >> it;
+            }
+        }
+        void print() {
+            for(auto &it: data) {
+                std::cout << it << " ";
+            }
+            std::cout << "\n";
+        }
+        void mul(multivar<base> const& b) {
+            assert(dim == b.dim);
+            size_t K = size(dim);
+            if(K == 0) {
+                data[0] *= b.data[0];
+                return;
+            }
+            std::vector<dft<base>> A, B;
+            size_t M = std::max(flen, std::bit_ceil(2 * N - 1) / 2);
+            for(size_t i = 0; i < K; i++) {
+                A.emplace_back(data | std::views::enumerate | std::views::transform(
+                    [&](auto jx) {
+                        auto [j, x] = jx;
+                        return ranks[j] == i ? x : base(0);
+                    }
+                ), M);
+                B.emplace_back(b.data | std::views::enumerate | std::views::transform(
+                    [&](auto jx) {
+                        auto [j, x] = jx;
+                        return ranks[j] == i ? x : base(0);
+                    }
+                ), M);
+            }
+            std::vector<dft<base>> C;
+            for(size_t i = 0; i < K; i++) {
+                dft<base> C(M);
+                cvector X = C.A;
+                for(size_t j = 0; j < K; j++) {
+                    size_t tj = (i - j + K) % K;
+                    A[j].template dot<false>(B[tj].A, B[tj].B, C.A, C.B, X);
+                }
+                std::vector<base, cp_algo::big_alloc<base>> res((N + flen - 1) / flen * flen);
+                C.A.ifft();
+                C.B.ifft();
+                X.ifft();
+                C.recover_mod(X, res, N);
+                for(size_t j = 0; j < N; j++) {
+                    if(i == ranks[j]) {
+                        data[j] = res[j];
+                    }
+                }
+            }
+        }
+    };
+}
+#endif // CP_ALGO_MATH_MULTIVAR_HPP

verify/poly/multivar.test.cpp

@@ -0,0 +1,32 @@
+// @brief Multidimensional Convolution (Truncated)
+#define PROBLEM "https://judge.yosupo.jp/problem/multivariate_convolution"
+#pragma GCC optimize("Ofast,unroll-loops")
+#define CP_ALGO_CHECKPOINT
+#include <bits/stdc++.h>
+#include "blazingio/blazingio.min.hpp"
+#include "cp-algo/math/multivar.hpp"
+
+using namespace std;
+using namespace cp_algo::math::fft;
+
+const int mod = 998244353;
+using base = cp_algo::math::modint<mod>;
+
+void solve() {
+    int k;
+    cin >> k;
+    vector<size_t> N(k);
+    for(auto &n: N) cin >> n;
+    multivar<base> a(N), b(N);
+    a.read();
+    b.read();
+    a.mul(b);
+    a.print();
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    solve();
+}