library/convolution/lcm-convolution.hpp

• View this file on GitHub
• Last update: 2023-12-03 13:41:56+08:00
• Include: #include "library/convolution/lcm-convolution.hpp"

Depends on

• library/math/divisor-multiple-transform.hpp
• library/math/prime-enumerate.hpp

Code

#pragma once
#include <vector>

#include <cassert>

#include "../math/divisor-multiple-transform.hpp"


namespace felix {

template<class T>
std::vector<T> gcd_convolution(const std::vector<T>& a, const std::vector<T>& b) {
	assert(a.size() == b.size());
	auto f = a, g = b;
	divisor_transform::zeta_transform(f);
	divisor_transform::zeta_transform(g);
	for(int i = 0; i < (int) f.size(); i++) {
		f[i] *= g[i];
	}
	divisor_transform::mobius_transform(f);
	return f;
}

} // namespace felix

#line 2 "library/convolution/lcm-convolution.hpp"
#include <vector>

#include <cassert>

#line 3 "library/math/prime-enumerate.hpp"
#include <cmath>


namespace felix {

// 2, 3, 5, 7, ...

std::vector<int> prime_enumerate(int N) {
	std::vector<bool> sieve(N / 3 + 1, 1);
	for(int p = 5, d = 4, i = 1, sqn = std::sqrt(N); p <= sqn; p += d = 6 - d, i++) {
		if(!sieve[i]) {
			continue;
		}
		for(int q = p * p / 3, r = d * p / 3 + (d * p % 3 == 2), s = 2 * p, qe = sieve.size(); q < qe; q += r = s - r) {
			sieve[q] = 0;
		}
	}
	std::vector<int> ret{2, 3};
	for(int p = 5, d = 4, i = 1; p <= N; p += d = 6 - d, i++) {
		if(sieve[i]) {
			ret.push_back(p);
		}
	}
	while(!ret.empty() && ret.back() > N) {
		ret.pop_back();
	}
	return ret;
}

} // namespace felix

#line 4 "library/math/divisor-multiple-transform.hpp"

namespace felix {

struct divisor_transform {
	template<class T>
	static void zeta_transform(std::vector<T>& a) {
		int n = a.size() - 1;
		for(auto p : prime_enumerate(n)) {
			for(int i = 1; i * p <= n; i++) {
				a[i * p] += a[i];
			}
		}
	}

	template<class T>
	static void mobius_transform(std::vector<T>& a) {
		int n = a.size() - 1;
		for(auto p : prime_enumerate(n)) {
			for(int i = n / p; i > 0; i--) {
				a[i * p] -= a[i];
			}
		}
	}
};

struct multiple_transform {
	template<class T>
	static void zeta_transform(std::vector<T>& a) {
		int n = a.size() - 1;
		for(auto p : prime_enumerate(n)) {
			for(int i = n / p; i > 0; i--) {
				a[i] += a[i * p];
			}
		}
	}

	template<class T>
	static void mobius_transform(std::vector<T>& a) {
		int n = a.size() - 1;
		for(auto p : prime_enumerate(n)) {
			for(int i = 1; i * p <= n; i++) {
				a[i] -= a[i * p];
			}
		}
	}
};

} // namespace felix

#line 5 "library/convolution/lcm-convolution.hpp"

namespace felix {

template<class T>
std::vector<T> gcd_convolution(const std::vector<T>& a, const std::vector<T>& b) {
	assert(a.size() == b.size());
	auto f = a, g = b;
	divisor_transform::zeta_transform(f);
	divisor_transform::zeta_transform(g);
	for(int i = 0; i < (int) f.size(); i++) {
		f[i] *= g[i];
	}
	divisor_transform::mobius_transform(f);
	return f;
}

} // namespace felix

Back to top page