Add and convolution and test

Author: adamant-pwn

cp-algo/math/and_convolution.hpp

@@ -0,0 +1,62 @@
+#ifndef CP_ALGO_MATH_AND_CONVOLUTION_HPP
+#define CP_ALGO_MATH_AND_CONVOLUTION_HPP
+#include "../number_theory/modint.hpp"
+#include "../util/bit.hpp"
+#include "../util/checkpoint.hpp"
+
+namespace cp_algo::math {
+    enum transform_dir { forward, inverse };
+    // Recursive Zeta / Möbius (AND) transform for size N (power of two)
+    template<auto N, transform_dir direction>
+    void and_transform(auto &&a) {
+        if constexpr (N == 1) {
+            return;
+        } else {
+            constexpr auto half = N / 2;
+            and_transform<half, direction>(&a[0]);
+            and_transform<half, direction>(&a[half]);
+            for (uint32_t i = 0; i < half; i++) {
+                if constexpr (direction == forward) {
+                    a[i] += a[i + half];
+                } else {
+                    a[i] -= a[i + half];
+                }
+            }
+        }
+    }
+    // Wrapper that deduces N at compile time via with_bit_floor
+    template<transform_dir direction>
+    inline void and_transform(auto &&a, auto n) {
+        with_bit_floor(n, [&]<auto NN>() {
+            assert(NN == n);
+            and_transform<NN, direction>(a);
+        });
+    }
+    template<transform_dir direction = forward>
+    inline void and_transform(auto &&a) {
+        and_transform<direction>(a, std::size(a));
+    }
+    // In-place AND convolution on sequences of equal length (power of two)
+    void and_convolution_inplace(auto &a, auto &b) {
+        auto N = static_cast<uint32_t>(std::size(a));
+        and_transform(a);
+        and_transform(b);
+        checkpoint("transform");
+        for (uint32_t i = 0; i < N; i++) {
+            a[i] *= b[i];
+        }
+        checkpoint("dot");
+        and_transform<inverse>(a);
+        checkpoint("itransform");
+    }
+
+    // Returns AND convolution of a and b; pads to next power of two
+    auto and_convolution(auto a, auto b) {
+        auto n = std::bit_ceil(std::max(std::size(a), std::size(b)));
+        a.resize(n);
+        b.resize(n);
+        and_convolution_inplace(a, b);
+        return a;
+    }
+}
+#endif // CP_ALGO_MATH_AND_CONVOLUTION_HPP

verify/math/and_convolution.test.cpp

@@ -0,0 +1,41 @@
+// @brief Bitwise And Convolution
+#define PROBLEM "https://judge.yosupo.jp/problem/bitwise_and_convolution"
+#pragma GCC optimize("O3,unroll-loops")
+#include <bits/allocator.h>
+#pragma GCC target("avx2")
+#include <iostream>
+#include "blazingio/blazingio.min.hpp"
+#define CP_ALGO_CHECKPOINT
+#include "cp-algo/number_theory/modint.hpp"
+#include "cp-algo/util/big_alloc.hpp"
+#include "cp-algo/util/checkpoint.hpp"
+#include "cp-algo/math/and_convolution.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+
+const int mod = 998244353;
+using base = cp_algo::math::modint<mod>;
+
+void solve() {
+    uint32_t n;
+    cin >> n;
+    uint32_t N = 1u << n;
+    cp_algo::big_vector<base> a(N), b(N);
+    for (auto &it : a) {cin >> it;}
+    for (auto &it : b) {cin >> it;}
+    cp_algo::checkpoint("read");
+    cp_algo::math::and_convolution_inplace(a, b);
+    for (auto it : a) {cout << it << ' ';}
+    cp_algo::checkpoint("write");
+    cp_algo::checkpoint<1>();
+}
+
+signed main() {
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    while (t--) {
+        solve();
+    }
+}