test/unit/zeta_mobius_transform.test.cpp

• View this file on GitHub
• Last update: 2026-10-27 06:53:58+09:00
• Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A

Depends on

• Bitwise And Convolution (math/bitwise_and_convolution.hpp)
• Bitwise Or Convolution (math/bitwise_or_convolution.hpp)
• 高速ゼータ変換・高速メビウス変換 (math/zeta_mobius_transform.hpp)

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <atcoder/modint>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#include "math/bitwise_and_convolution.hpp"
#include "math/bitwise_or_convolution.hpp"
#include "math/zeta_mobius_transform.hpp"
using mint = atcoder::modint998244353;
int main() {
    std::mt19937 rng(7959);
    for (int k = 0; k <= 8; k++) {
        int n = 1 << k;
        for (int trial = 0; trial < 5; trial++) {
            std::vector<mint> f(n);
            for (auto& x : f) x = rng();
            std::vector<mint> sup(n), sub(n);
            for (int s = 0; s < n; s++) {
                for (int t = 0; t < n; t++) {
                    if ((s & t) == s) sup[s] += f[t];
                    if ((s & t) == t) sub[s] += f[t];
                }
            }
            auto g = f;
            superset_zeta(g);
            assert(g == sup);
            superset_mobius(g);
            assert(g == f);
            g = f;
            subset_zeta(g);
            assert(g == sub);
            subset_mobius(g);
            assert(g == f);
        }
    }
    // long long: convolutions against brute force, including n = 1
    for (int k = 0; k <= 6; k++) {
        int n = 1 << k;
        for (int trial = 0; trial < 5; trial++) {
            std::vector<long long> a(n), b(n), ca(n), co(n);
            for (auto& x : a) x = int(rng() % 2001) - 1000;
            for (auto& x : b) x = int(rng() % 2001) - 1000;
            for (int i = 0; i < n; i++) {
                for (int j = 0; j < n; j++) {
                    ca[i & j] += a[i] * b[j];
                    co[i | j] += a[i] * b[j];
                }
            }
            assert(bitwise_and_convolution(a, b) == ca);
            assert(bitwise_or_convolution(a, b) == co);
        }
    }
    std::cout << "Hello World" << std::endl;
    return 0;
}

#line 1 "test/unit/zeta_mobius_transform.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <atcoder/modint>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#line 4 "math/zeta_mobius_transform.hpp"
template <class T>
void superset_zeta(std::vector<T>& f) {
    int n = int(f.size());
    assert((n & -n) == n);
    for (int i = 1; i < n; i <<= 1) {
        for (int j = 0; j < n; j++) {
            if ((j & i) == 0) {
                f[j] += f[j | i];
            }
        }
    }
}
template <class T>
void superset_mobius(std::vector<T>& f) {
    int n = int(f.size());
    assert((n & -n) == n);
    for (int i = 1; i < n; i <<= 1) {
        for (int j = 0; j < n; j++) {
            if ((j & i) == 0) {
                f[j] -= f[j | i];
            }
        }
    }
}
template <class T>
void subset_zeta(std::vector<T>& f) {
    int n = int(f.size());
    assert((n & -n) == n);
    for (int i = 1; i < n; i <<= 1) {
        for (int j = 0; j < n; j++) {
            if ((j & i) == 0) {
                f[j | i] += f[j];
            }
        }
    }
}
template <class T>
void subset_mobius(std::vector<T>& f) {
    int n = int(f.size());
    assert((n & -n) == n);
    for (int i = 1; i < n; i <<= 1) {
        for (int j = 0; j < n; j++) {
            if ((j & i) == 0) {
                f[j | i] -= f[j];
            }
        }
    }
}
#line 3 "math/bitwise_and_convolution.hpp"
template <class T>
std::vector<T> bitwise_and_convolution(std::vector<T> a,
                                          std::vector<T> b) {
    superset_zeta(a);
    superset_zeta(b);
    for (int i = 0; i < int(a.size()); i++) {
        a[i] *= b[i];
    }
    superset_mobius(a);
    return a;
}
#line 3 "math/bitwise_or_convolution.hpp"
template <class T>
std::vector<T> bitwise_or_convolution(std::vector<T> a,
                                         std::vector<T> b) {
    subset_zeta(a);
    subset_zeta(b);
    for (int i = 0; i < int(a.size()); i++) {
        a[i] *= b[i];
    }
    subset_mobius(a);
    return a;
}
#line 11 "test/unit/zeta_mobius_transform.test.cpp"
using mint = atcoder::modint998244353;
int main() {
    std::mt19937 rng(7959);
    for (int k = 0; k <= 8; k++) {
        int n = 1 << k;
        for (int trial = 0; trial < 5; trial++) {
            std::vector<mint> f(n);
            for (auto& x : f) x = rng();
            std::vector<mint> sup(n), sub(n);
            for (int s = 0; s < n; s++) {
                for (int t = 0; t < n; t++) {
                    if ((s & t) == s) sup[s] += f[t];
                    if ((s & t) == t) sub[s] += f[t];
                }
            }
            auto g = f;
            superset_zeta(g);
            assert(g == sup);
            superset_mobius(g);
            assert(g == f);
            g = f;
            subset_zeta(g);
            assert(g == sub);
            subset_mobius(g);
            assert(g == f);
        }
    }
    // long long: convolutions against brute force, including n = 1
    for (int k = 0; k <= 6; k++) {
        int n = 1 << k;
        for (int trial = 0; trial < 5; trial++) {
            std::vector<long long> a(n), b(n), ca(n), co(n);
            for (auto& x : a) x = int(rng() % 2001) - 1000;
            for (auto& x : b) x = int(rng() % 2001) - 1000;
            for (int i = 0; i < n; i++) {
                for (int j = 0; j < n; j++) {
                    ca[i & j] += a[i] * b[j];
                    co[i | j] += a[i] * b[j];
                }
            }
            assert(bitwise_and_convolution(a, b) == ca);
            assert(bitwise_or_convolution(a, b) == co);
        }
    }
    std::cout << "Hello World" << std::endl;
    return 0;
}

Back to top page