bitwise_conv.hpp

• View this file on GitHub
• Last update: 2025-01-19 15:28:01+08:00
• Include: #include "bitwise_conv.hpp"

Depends on

• sps_fft.hpp
• subset_conv.hpp

Verified with

• test/convolution/bitwise_and_convolution.0.test.cpp
• test/convolution/bitwise_xor_convolution.0.test.cpp

Code

#pragma once

#include "sps_fft.hpp"
#include "subset_conv.hpp"
#include <algorithm>
#include <cassert>
#include <vector>

template<typename Tp>
inline std::vector<Tp> bitwise_or_convolution(std::vector<Tp> a, std::vector<Tp> b) {
    assert(a.size() == b.size());
    const int n = a.size();
    assert((n & (n - 1)) == 0);
    assert(n > 0);
    subset_zeta(a);
    subset_zeta(b);
    for (int i = 0; i < n; ++i) a[i] *= b[i];
    subset_moebius(a);
    return a;
}

template<typename Tp>
inline std::vector<Tp> bitwise_and_convolution(const std::vector<Tp> &a, const std::vector<Tp> &b) {
    auto ab = bitwise_or_convolution(std::vector(a.rbegin(), a.rend()),
                                     std::vector(b.rbegin(), b.rend()));
    std::reverse(ab.begin(), ab.end());
    return ab;
}

template<typename Tp>
inline std::vector<Tp> bitwise_xor_convolution(std::vector<Tp> a, std::vector<Tp> b) {
    assert(a.size() == b.size());
    const int n = a.size();
    assert((n & (n - 1)) == 0);
    assert(n > 0);
    sps_fft(a);
    sps_fft(b);
    for (int i = 0; i < n; ++i) a[i] *= b[i];
    sps_inv_fft(a);
    return a;
}

#line 2 "bitwise_conv.hpp"

#line 2 "sps_fft.hpp"

#include <cassert>
#include <iterator>
#include <vector>

// set power series = R[x_1,...,x_n]/(x_1^2,...,x_n^2)
// FFT is computing F({1,-1}^n)

template<typename Iterator> inline void sps_fft_n(Iterator a, int n) {
    assert((n & (n - 1)) == 0);
    for (int i = 2; i <= n; i *= 2)
        for (int j = 0; j < n; j += i)
            for (int k = j; k < j + i / 2; ++k) {
                const auto u = a[k], v = a[k + i / 2];
                a[k] = u + v, a[k + i / 2] = u - v;
            }
}

template<typename Tp> inline void sps_fft(std::vector<Tp> &a) { sps_fft_n(a.begin(), a.size()); }

template<typename Iterator> inline void sps_inv_fft_n(Iterator a, int n) {
    using Tp = typename std::iterator_traits<Iterator>::value_type;
    sps_fft_n(a, n);
    const Tp iv = Tp::mod() - (Tp::mod() - 1) / n;
    for (int i = 0; i < n; ++i) a[i] *= iv;
}

template<typename Tp> inline void sps_inv_fft(std::vector<Tp> &a) {
    sps_inv_fft_n(a.begin(), a.size());
}
#line 2 "subset_conv.hpp"

#line 5 "subset_conv.hpp"

template<typename Tp> inline std::vector<std::vector<Tp>> to_ranked(const std::vector<Tp> &A) {
    const int N = A.size();
    int LogN    = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector res(LogN + 1, std::vector<Tp>(N));
    for (int i = 0; i < N; ++i) res[__builtin_popcount(i)][i] = A[i];
    return res;
}

template<typename Tp> inline std::vector<Tp> from_ranked(const std::vector<std::vector<Tp>> &A) {
    const int N = A[0].size();
    std::vector<Tp> res(N);
    for (int i = 0; i < N; ++i) res[i] = A[__builtin_popcount(i)][i];
    return res;
}

template<typename Iterator> inline void subset_zeta_n(Iterator a, int n) {
    assert((n & (n - 1)) == 0);
    for (int i = 2; i <= n; i *= 2)
        for (int j = 0; j < n; j += i)
            for (int k = j; k < j + i / 2; ++k) a[k + i / 2] += a[k];
}

template<typename Tp> inline void subset_zeta(std::vector<Tp> &a) {
    subset_zeta_n(a.begin(), a.size());
}

template<typename Iterator> inline void subset_moebius_n(Iterator a, int n) {
    assert((n & (n - 1)) == 0);
    for (int i = 2; i <= n; i *= 2)
        for (int j = 0; j < n; j += i)
            for (int k = j; k < j + i / 2; ++k) a[k + i / 2] -= a[k];
}

template<typename Tp> inline void subset_moebius(std::vector<Tp> &a) {
    subset_moebius_n(a.begin(), a.size());
}

template<typename Tp>
inline std::vector<Tp> subset_convolution(const std::vector<Tp> &A, const std::vector<Tp> &B) {
    assert(A.size() == B.size());
    const int N = A.size();
    int LogN    = 0;
    while ((1 << LogN) != N) ++LogN;
    auto rankedA = to_ranked(A);
    auto rankedB = to_ranked(B);

    for (int i = 0; i <= LogN; ++i) {
        subset_zeta(rankedA[i]);
        subset_zeta(rankedB[i]);
    }

    // see: https://codeforces.com/blog/entry/126418
    // see: https://oeis.org/A025480
    std::vector<int> map(LogN + 1);
    for (int i = 0; i <= LogN; ++i) map[i] = (i & 1) ? map[i / 2] : i / 2;

    std::vector rankedAB(LogN / 2 + 1, std::vector<Tp>(N));
    for (int i = 0; i <= LogN; ++i)
        for (int j = 0; i + j <= LogN; ++j)
            for (int k = (1 << j) - 1; k < N; ++k)
                rankedAB[map[i + j]][k] += rankedA[i][k] * rankedB[j][k];

    for (int i = 0; i <= LogN / 2; ++i) subset_moebius(rankedAB[i]);

    std::vector<Tp> res(N);
    for (int i = 0; i < N; ++i) res[i] = rankedAB[map[__builtin_popcount(i)]][i];
    return res;
}
#line 5 "bitwise_conv.hpp"
#include <algorithm>
#line 8 "bitwise_conv.hpp"

template<typename Tp>
inline std::vector<Tp> bitwise_or_convolution(std::vector<Tp> a, std::vector<Tp> b) {
    assert(a.size() == b.size());
    const int n = a.size();
    assert((n & (n - 1)) == 0);
    assert(n > 0);
    subset_zeta(a);
    subset_zeta(b);
    for (int i = 0; i < n; ++i) a[i] *= b[i];
    subset_moebius(a);
    return a;
}

template<typename Tp>
inline std::vector<Tp> bitwise_and_convolution(const std::vector<Tp> &a, const std::vector<Tp> &b) {
    auto ab = bitwise_or_convolution(std::vector(a.rbegin(), a.rend()),
                                     std::vector(b.rbegin(), b.rend()));
    std::reverse(ab.begin(), ab.end());
    return ab;
}

template<typename Tp>
inline std::vector<Tp> bitwise_xor_convolution(std::vector<Tp> a, std::vector<Tp> b) {
    assert(a.size() == b.size());
    const int n = a.size();
    assert((n & (n - 1)) == 0);
    assert(n > 0);
    sps_fft(a);
    sps_fft(b);
    for (int i = 0; i < n; ++i) a[i] *= b[i];
    sps_inv_fft(a);
    return a;
}

Back to top page